multifields support

carlosdelest · Mar 22, 2024 · c80677f · c80677f
1 parent 05aa06f
commit c80677f
Show file tree

Hide file tree

Showing 3 changed files with 72 additions and 14 deletions.
diff --git a/.../java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilter.java b/.../java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilter.java
@@ -291,8 +291,9 @@ private Map<String, List<FieldInferenceRequest>> createFieldInferenceRequests(Bu
                 final Map<String, Object> docMap = indexRequest.sourceAsMap();
                 for (var entry : fieldInferenceMetadata.getFieldInferenceOptions().entrySet()) {
                     String inferenceId = entry.getValue().inferenceId();
-                    for (var field : entry.getValue().sourceFields()) {
-                        var value = XContentMapValues.extractValue(field, docMap);
+                    String fieldName = entry.getKey();
+                    for (var sourceField : entry.getValue().sourceFields()) {
+                        var value = XContentMapValues.extractValue(sourceField, docMap);
                         if (value == null) {
                             continue;
                         }
@@ -311,13 +312,13 @@ private Map<String, List<FieldInferenceRequest>> createFieldInferenceRequests(Bu
                                 inferenceId,
                                 k -> new ArrayList<>()
                             );
-                            fieldRequests.add(new FieldInferenceRequest(item.id(), field, valueStr));
+                            fieldRequests.add(new FieldInferenceRequest(item.id(), fieldName, valueStr));
                         } else {
                             inferenceResults.get(item.id()).failures.add(
                                 new ElasticsearchStatusException(
                                     "Invalid format for field [{}], expected [String] got [{}]",
                                     RestStatus.BAD_REQUEST,
-                                    field,
+                                    fieldName,
                                     value.getClass().getSimpleName()
                                 )
                             );

diff --git a/...ce/src/main/java/org/elasticsearch/xpack/inference/mapper/InferenceResultFieldMapper.java b/...ce/src/main/java/org/elasticsearch/xpack/inference/mapper/InferenceResultFieldMapper.java
@@ -148,10 +148,15 @@ public InferenceResultFieldMapper() {
 
     @Override
     protected void parseCreateField(DocumentParserContext context) throws IOException {
-        XContentParser parser = context.parser();
-        failIfTokenIsNot(parser, XContentParser.Token.START_OBJECT);
-
-        parseAllFields(context);
+        boolean withinLeafObject = context.path().isWithinLeafObject();
+        try {
+            context.path().setWithinLeafObject(true);
+            XContentParser parser = context.parser();
+            failIfTokenIsNot(parser, XContentParser.Token.START_OBJECT);
+            parseAllFields(context);
+        } finally {
+            context.path().setWithinLeafObject(withinLeafObject);
+        }
     }
 
     private static void parseAllFields(DocumentParserContext context) throws IOException {
@@ -164,12 +169,18 @@ private static void parseAllFields(DocumentParserContext context) throws IOExcep
         }
     }
 
-    private static void parseSingleField(DocumentParserContext context, MapperBuilderContext mapperBuilderContext) throws IOException {
+    private static void parseSingleField(DocumentParserContext context, MapperBuilderContext mapperBuilderContext)
+        throws IOException {
 
         XContentParser parser = context.parser();
         String fieldName = parser.currentName();
-        Mapper mapper = context.getMapper(fieldName);
-        if (mapper == null || SemanticTextFieldMapper.CONTENT_TYPE.equals(mapper.typeName()) == false) {
+        Mapper mapper = findMapper(context.root(), fieldName);
+        if (mapper == null) {
+            throw new DocumentParsingException(
+                parser.getTokenLocation(),
+                Strings.format("Field [%s] is not registered as a field type", fieldName)
+            );
+        } else if (SemanticTextFieldMapper.CONTENT_TYPE.equals(mapper.typeName()) == false)  {
             throw new DocumentParsingException(
                 parser.getTokenLocation(),
                 Strings.format("Field [%s] is not registered as a %s field type", fieldName, SemanticTextFieldMapper.CONTENT_TYPE)
@@ -190,18 +201,28 @@ private static void parseSingleField(DocumentParserContext context, MapperBuilde
                     fieldName,
                     modelSettings
                 );
-                parseFieldInferenceChunks(context, mapperBuilderContext, fieldName, modelSettings, nestedObjectMapper);
+                parseFieldInferenceChunks(context, modelSettings, nestedObjectMapper);
             } else {
                 logger.debug("Skipping unrecognized field name [" + currentName + "]");
                 advancePastCurrentFieldName(parser);
             }
         }
     }
 
+    private static Mapper findMapper(Mapper mapper, String fullPath) {
+        String[] pathElements = fullPath.split("\\.");
+        for (int i = 0; i < pathElements.length; i++) {
+            Mapper next = mapper.getMapper(pathElements[i]);
+            if (next == null) {
+                return null;
+            }
+            mapper = next;
+        }
+        return mapper;
+    }
+
     private static void parseFieldInferenceChunks(
         DocumentParserContext context,
-        MapperBuilderContext mapperBuilderContext,
-        String fieldName,
         SemanticTextModelSettings modelSettings,
         NestedObjectMapper nestedObjectMapper
     ) throws IOException {
@@ -243,6 +264,8 @@ private static void parseFieldInferenceChunkElement(
             if (childMapper instanceof FieldMapper fieldMapper) {
                 parser.nextToken();
                 fieldMapper.parse(childContext);
+                // Reset leaf object after parsing the field
+                context.path().setWithinLeafObject(true);
             } else {
                 // This should never happen, but fail parsing if it does so that it's not a silent failure
                 throw new DocumentParsingException(

diff --git a/...ce/src/yamlRestTest/resources/rest-api-spec/test/inference/10_semantic_text_inference.yml b/...ce/src/yamlRestTest/resources/rest-api-spec/test/inference/10_semantic_text_inference.yml
@@ -309,3 +309,37 @@ setup:
         id: doc_1
         body:
           non_inference_field: "non inference test"
+
+
+---
+"semantic_text multifields calculate inference for parent field":
+  - do:
+      indices.create:
+        index: test-multifield-index
+        body:
+          mappings:
+            properties:
+              top_level_field:
+                type: text
+                fields:
+                  semantic_multifield:
+                    type: semantic_text
+                    model_id: dense-inference-id
+  - do:
+      index:
+        index: test-multifield-index
+        id: doc_1
+        body:
+          top_level_field: "multifield inference test"
+
+  - do:
+      get:
+        index: test-multifield-index
+        id: doc_1
+
+  - match: { _source.top_level_field: "multifield inference test" }
+  - match: { _source._inference.top_level_field\.semantic_multifield.results.0.text: "multifield inference test" }
+  - exists: _source._inference.top_level_field\.semantic_multifield.results.0.inference
+  - exists: _source._inference.top_level_field\.semantic_multifield.results.0.inference
+
+