From 3028bd6e250a663fcacfe0479c1ee3c8b75e5e30 Mon Sep 17 00:00:00 2001
From: carlosdelest <carlos.delgado@elastic.co>
Date: Mon, 20 Nov 2023 18:18:48 +0100
Subject: [PATCH] Add serialization of models for fields in toXContent so load
 IndexMetadata from disk works

---
 .../cluster/metadata/IndexMetadata.java            | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexMetadata.java b/server/src/main/java/org/elasticsearch/cluster/metadata/IndexMetadata.java
index 5b62c3eb17adf..c6f496eeba9b9 100644
--- a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexMetadata.java
+++ b/server/src/main/java/org/elasticsearch/cluster/metadata/IndexMetadata.java
@@ -79,6 +79,7 @@
 import java.util.OptionalLong;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 import static org.elasticsearch.TransportVersions.SEMANTIC_TEXT_FIELD;
 import static org.elasticsearch.cluster.metadata.Metadata.CONTEXT_MODE_PARAM;
@@ -546,6 +547,7 @@ public Iterator<Setting<?>> settings() {
     public static final String KEY_WRITE_LOAD_FORECAST = "write_load_forecast";
 
     public static final String KEY_SHARD_SIZE_FORECAST = "shard_size_forecast";
+    public static final String INFERENCE_MODELS_FIELDS = "inference_models_fields";
 
     public static final String INDEX_STATE_FILE_PREFIX = "state-";
 
@@ -2405,6 +2407,11 @@ public static void toXContent(IndexMetadata indexMetadata, XContentBuilder build
                 builder.field(KEY_SHARD_SIZE_FORECAST, indexMetadata.shardSizeInBytesForecast);
             }
 
+            Map<String, List<String>> inferenceModelsForFields = indexMetadata.getInferenceModelsForFields();
+            if ((inferenceModelsForFields != null) && (inferenceModelsForFields.isEmpty() == false)) {
+                builder.field(INFERENCE_MODELS_FIELDS, indexMetadata.getInferenceModelsForFields());
+            }
+
             builder.endObject();
         }
 
@@ -2482,6 +2489,13 @@ public static IndexMetadata fromXContent(XContentParser parser, Map<String, Mapp
                         case KEY_STATS:
                             builder.stats(IndexMetadataStats.fromXContent(parser));
                             break;
+                        case INFERENCE_MODELS_FIELDS:
+                            Map<String, List<String>> inferenceModels = parser.map(HashMap::new, XContentParser::list)
+                                .entrySet()
+                                .stream()
+                                .collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().stream().map(Object::toString).toList()));
+                            builder.inferenceModelsForfields(inferenceModels);
+                            break;
                         default:
                             // assume it's custom index metadata
                             builder.putCustom(currentFieldName, parser.mapStrings());