Some minor changes for field update

carlosdelest · Oct 26, 2023 · ddc966b · ddc966b
1 parent 6da59a3
commit ddc966b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 31 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/SemanticTextFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/SemanticTextFieldMapper.java
@@ -24,10 +24,9 @@
 public class SemanticTextFieldMapper extends FieldMapper {
 
     public static final String CONTENT_TYPE = "semantic_text";
-    private static final String SPARSE_VECTOR_SUFFIX = "_inference";
 
-    private static final String TEXT_SUBFIELD_NAME = "text";
-    private static final String SPARSE_VECTOR_SUBFIELD_NAME = "inference";
+    public static final String TEXT_SUBFIELD_NAME = "text";
+    public static final String SPARSE_VECTOR_SUBFIELD_NAME = "inference";
 
     private static SemanticTextFieldMapper toType(FieldMapper in) {
         return (SemanticTextFieldMapper) in;
@@ -55,33 +54,25 @@ public Builder modelId(String modelId) {
 
         @Override
         protected Parameter<?>[] getParameters() {
-            return new Parameter<?>[] {
-                modelId,
-                meta };
-        }
-
-        private SemanticTextFieldType buildFieldType(
-            MapperBuilderContext context
-        ) {
-                return new SemanticTextFieldType(
-                    context.buildFullName(name),
-                    modelId.getValue(),
-                    meta.getValue()
-            );
+            return new Parameter<?>[] { modelId, meta };
+        }
+
+        private SemanticTextFieldType buildFieldType(MapperBuilderContext context) {
+            return new SemanticTextFieldType(context.buildFullName(name), modelId.getValue(), meta.getValue());
         }
 
         @Override
         public SemanticTextFieldMapper build(MapperBuilderContext context) {
             SemanticTextFieldType stft = new SemanticTextFieldType(context.buildFullName(name), modelId.getValue(), meta.getValue());
-            String fieldName = name() + SPARSE_VECTOR_SUFFIX;
-            SubFieldInfo sparseVectorFieldInfo = new SubFieldInfo(fieldName, new SparseVectorFieldMapper.Builder(fieldName).build(context));
+            SubFieldInfo sparseVectorFieldInfo = new SubFieldInfo(
+                SPARSE_VECTOR_SUBFIELD_NAME,
+                new SparseVectorFieldMapper.Builder(SPARSE_VECTOR_SUBFIELD_NAME).build(context)
+            );
             return new SemanticTextFieldMapper(name, stft, modelId.getValue(), sparseVectorFieldInfo, copyTo, this);
         }
     }
 
-    public static final TypeParser PARSER = new TypeParser(
-        (n, c) ->  new Builder(n), notInMultiFields(CONTENT_TYPE)
-    );
+    public static final TypeParser PARSER = new TypeParser((n, c) -> new Builder(n), notInMultiFields(CONTENT_TYPE));
 
     private static final class SubFieldInfo {
 
@@ -101,13 +92,9 @@ public static class SemanticTextFieldType extends SimpleMappedFieldType {
 
         private final String modelId;
 
-        public SemanticTextFieldType(
-            String name,
-            String modelId,
-            Map<String, String> meta
-        ) {
+        public SemanticTextFieldType(String name, String modelId, Map<String, String> meta) {
             super(name, true, false, false, TextSearchInfo.NONE, meta);
-            this.modelId =  modelId;
+            this.modelId = modelId;
         }
 
         public String modelId() {
@@ -172,7 +159,8 @@ public void parse(DocumentParserContext context) throws IOException {
 
         boolean textFound = false;
         boolean inferenceFound = false;
-        for (XContentParser.Token token = context.parser().nextToken(); token != XContentParser.Token.END_OBJECT; token = context.parser().nextToken()) {
+        for (XContentParser.Token token = context.parser().nextToken(); token != XContentParser.Token.END_OBJECT; token = context.parser()
+            .nextToken()) {
             if (token != XContentParser.Token.FIELD_NAME) {
                 throw new IllegalArgumentException("[semantic_text] fields expect an object with field names, found " + token);
             }

diff --git a/server/src/main/java/org/elasticsearch/ingest/FieldInferenceBulkRequestPreprocessor.java b/server/src/main/java/org/elasticsearch/ingest/FieldInferenceBulkRequestPreprocessor.java
@@ -16,9 +16,11 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.client.internal.OriginSettingClient;
 import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.index.mapper.SemanticTextFieldMapper;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.plugins.internal.DocumentParsingObserver;
 
+import java.util.HashMap;
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.IntConsumer;
@@ -105,10 +107,12 @@ private void runInferenceForField(
         client.execute(InferenceAction.INSTANCE, inferenceRequest, ActionListener.runAfter(new ActionListener<InferenceAction.Response>() {
             @Override
             public void onResponse(InferenceAction.Response response) {
-                ingestDocument.removeField(fieldName);
                 // Transform into two subfields, one with the actual text and other with the inference
-                ingestDocument.setFieldValue(fieldName + "._text", fieldValue);
-                ingestDocument.setFieldValue(fieldName + "._inference", response.getResult().asMap(fieldName).get(fieldName));
+                Map<String, Object> newFieldValue = new HashMap<>();
+                newFieldValue.put(SemanticTextFieldMapper.TEXT_SUBFIELD_NAME, fieldValue);
+                newFieldValue.put(SemanticTextFieldMapper.SPARSE_VECTOR_SUBFIELD_NAME, response.getResult().asMap(fieldName).get(fieldName));
+                ingestDocument.setFieldValue(fieldName, newFieldValue);
+
                 updateIndexRequestSource(indexRequest, ingestDocument);
             }