First iteration for adding inference support for copy_to / multifields

carlosdelest · Mar 19, 2024 · d3f9d86 · d3f9d86
1 parent f029015
commit d3f9d86
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 31 deletions.
diff --git a/.../java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilter.java b/.../java/org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilter.java
@@ -290,34 +290,38 @@ private Map<String, List<FieldInferenceRequest>> createFieldInferenceRequests(Bu
                 }
                 final Map<String, Object> docMap = indexRequest.sourceAsMap();
                 for (var entry : fieldInferenceMetadata.getFieldInferenceOptions().entrySet()) {
-                    String field = entry.getKey();
                     String inferenceId = entry.getValue().inferenceId();
-                    var value = XContentMapValues.extractValue(field, docMap);
-                    if (value == null) {
-                        continue;
-                    }
-                    if (inferenceResults.get(item.id()) == null) {
-                        inferenceResults.set(
-                            item.id(),
-                            new FieldInferenceResponseAccumulator(
+                    for (var field : entry.getValue().sourceFields()) {
+                        var value = XContentMapValues.extractValue(field, docMap);
+                        if (value == null) {
+                            continue;
+                        }
+                        if (inferenceResults.get(item.id()) == null) {
+                            inferenceResults.set(
                                 item.id(),
-                                Collections.synchronizedList(new ArrayList<>()),
-                                Collections.synchronizedList(new ArrayList<>())
-                            )
-                        );
-                    }
-                    if (value instanceof String valueStr) {
-                        List<FieldInferenceRequest> fieldRequests = fieldRequestsMap.computeIfAbsent(inferenceId, k -> new ArrayList<>());
-                        fieldRequests.add(new FieldInferenceRequest(item.id(), field, valueStr));
-                    } else {
-                        inferenceResults.get(item.id()).failures.add(
-                            new ElasticsearchStatusException(
-                                "Invalid format for field [{}], expected [String] got [{}]",
-                                RestStatus.BAD_REQUEST,
-                                field,
-                                value.getClass().getSimpleName()
-                            )
-                        );
+                                new FieldInferenceResponseAccumulator(
+                                    item.id(),
+                                    Collections.synchronizedList(new ArrayList<>()),
+                                    Collections.synchronizedList(new ArrayList<>())
+                                )
+                            );
+                        }
+                        if (value instanceof String valueStr) {
+                            List<FieldInferenceRequest> fieldRequests = fieldRequestsMap.computeIfAbsent(
+                                inferenceId,
+                                k -> new ArrayList<>()
+                            );
+                            fieldRequests.add(new FieldInferenceRequest(item.id(), field, valueStr));
+                        } else {
+                            inferenceResults.get(item.id()).failures.add(
+                                new ElasticsearchStatusException(
+                                    "Invalid format for field [{}], expected [String] got [{}]",
+                                    RestStatus.BAD_REQUEST,
+                                    field,
+                                    value.getClass().getSimpleName()
+                                )
+                            );
+                        }
                     }
                 }
             }

diff --git a/.../org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilterTests.java b/.../org/elasticsearch/xpack/inference/action/filter/ShardBulkInferenceActionFilterTests.java
@@ -92,7 +92,7 @@ public void testFilterNoop() throws Exception {
             new BulkItemRequest[0]
         );
         request.setFieldInferenceMetadata(
-            new FieldInferenceMetadata(Map.of("foo", new FieldInferenceMetadata.FieldInferenceOptions("bar", Set.of())))
+            new FieldInferenceMetadata(Map.of("foo", new FieldInferenceMetadata.FieldInferenceOptions("bar", Set.of("foo"))))
         );
         filter.apply(task, TransportShardBulkAction.ACTION_NAME, request, actionListener, actionFilterChain);
         awaitLatch(chainExecuted, 10, TimeUnit.SECONDS);
@@ -123,11 +123,11 @@ public void testInferenceNotFound() throws Exception {
         FieldInferenceMetadata inferenceFields = new FieldInferenceMetadata(
             Map.of(
                 "field1",
-                new FieldInferenceMetadata.FieldInferenceOptions(model.getInferenceEntityId(), Set.of()),
+                new FieldInferenceMetadata.FieldInferenceOptions(model.getInferenceEntityId(), Set.of("field1")),
                 "field2",
-                new FieldInferenceMetadata.FieldInferenceOptions("inference_0", Set.of()),
+                new FieldInferenceMetadata.FieldInferenceOptions("inference_0", Set.of("field2")),
                 "field3",
-                new FieldInferenceMetadata.FieldInferenceOptions("inference_0", Set.of())
+                new FieldInferenceMetadata.FieldInferenceOptions("inference_0", Set.of("field3"))
             )
         );
         BulkItemRequest[] items = new BulkItemRequest[10];
@@ -154,7 +154,7 @@ public void testManyRandomDocs() throws Exception {
         for (int i = 0; i < numInferenceFields; i++) {
             String field = randomAlphaOfLengthBetween(5, 10);
             String inferenceId = randomFrom(inferenceModelMap.keySet());
-            inferenceFieldsMap.put(field, new FieldInferenceMetadata.FieldInferenceOptions(inferenceId, Set.of()));
+            inferenceFieldsMap.put(field, new FieldInferenceMetadata.FieldInferenceOptions(inferenceId, Set.of(field)));
         }
         FieldInferenceMetadata fieldInferenceMetadata = new FieldInferenceMetadata(inferenceFieldsMap);