implement batch document update scenario for text embedding processor (#1217)

Signed-off-by: Will Hwang <[email protected]>

Author: will-hwang

src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java

@@ -6,17 +6,18 @@
 
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
 
 import org.opensearch.action.get.GetAction;
 import org.opensearch.action.get.GetRequest;
+import org.opensearch.action.get.MultiGetAction;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.core.common.util.CollectionUtils;
 import org.opensearch.env.Environment;
 import org.opensearch.ingest.IngestDocument;
+import org.opensearch.ingest.IngestDocumentWrapper;
 import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor;
 
 import lombok.extern.log4j.Log4j2;
@@ -74,29 +75,16 @@ public void doExecute(
         // if skipExisting flag is turned on, eligible inference texts will be compared and filtered after embeddings are copied
         String index = ingestDocument.getSourceAndMetadata().get(INDEX_FIELD).toString();
         String id = ingestDocument.getSourceAndMetadata().get(ID_FIELD).toString();
-        openSearchClient.execute(GetAction.INSTANCE, new GetRequest(index, id), ActionListener.wrap(response -> {
-            final Map<String, Object> existingDocument = response.getSourceAsMap();
-            if (existingDocument == null || existingDocument.isEmpty()) {
-                makeInferenceCall(ingestDocument, processMap, inferenceList, handler);
-                return;
-            }
-            // filter given ProcessMap by comparing existing document with ingestDocument
-            Map<String, Object> filteredProcessMap = textEmbeddingInferenceFilter.filter(
-                existingDocument,
-                ingestDocument.getSourceAndMetadata(),
-                processMap
-            );
-            // create inference list based on filtered ProcessMap
-            List<String> filteredInferenceList = createInferenceList(filteredProcessMap).stream()
-                .filter(Objects::nonNull)
-                .collect(Collectors.toList());
-            if (filteredInferenceList.isEmpty()) {
-                handler.accept(ingestDocument, null);
-            } else {
-                makeInferenceCall(ingestDocument, filteredProcessMap, filteredInferenceList, handler);
-            }
-
-        }, e -> { handler.accept(null, e); }));
+        openSearchClient.execute(
+            GetAction.INSTANCE,
+            new GetRequest(index, id),
+            ActionListener.wrap(
+                response -> getResponseHandler(response, ingestDocument, processMap, inferenceList, handler, textEmbeddingInferenceFilter),
+                e -> {
+                    handler.accept(null, e);
+                }
+            )
+        );
     }
 
     @Override
@@ -106,4 +94,47 @@ public void doBatchExecute(List<String> inferenceList, Consumer<List<?>> handler
             ActionListener.wrap(handler::accept, onException)
         );
     }
+
+    @Override
+    public void subBatchExecute(List<IngestDocumentWrapper> ingestDocumentWrappers, Consumer<List<IngestDocumentWrapper>> handler) {
+        try {
+            if (CollectionUtils.isEmpty(ingestDocumentWrappers)) {
+                handler.accept(ingestDocumentWrappers);
+                return;
+            }
+            List<DataForInference> dataForInferences = getDataForInference(ingestDocumentWrappers);
+            List<String> inferenceList = constructInferenceTexts(dataForInferences);
+            if (inferenceList.isEmpty()) {
+                handler.accept(ingestDocumentWrappers);
+                return;
+            }
+            // skip existing flag is turned off. Call doSubBatchExecute without filtering
+            if (skipExisting == false) {
+                doSubBatchExecute(ingestDocumentWrappers, inferenceList, dataForInferences, handler);
+                return;
+            }
+            // skipExisting flag is turned on, eligible inference texts in dataForInferences will be compared and filtered after embeddings
+            // are copied
+            openSearchClient.execute(
+                MultiGetAction.INSTANCE,
+                buildMultiGetRequest(dataForInferences),
+                ActionListener.wrap(
+                    response -> multiGetResponseHandler(
+                        response,
+                        ingestDocumentWrappers,
+                        inferenceList,
+                        dataForInferences,
+                        handler,
+                        textEmbeddingInferenceFilter
+                    ),
+                    e -> {
+                        // When exception is thrown in for MultiGetAction, set exception to all ingestDocumentWrappers
+                        updateWithExceptions(getIngestDocumentWrappers(dataForInferences), handler, e);
+                    }
+                )
+            );
+        } catch (Exception e) {
+            updateWithExceptions(ingestDocumentWrappers, handler, e);
+        }
+    }
 }

src/main/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessor.java

@@ -13,6 +13,7 @@
 import java.util.List;
 import java.util.ListIterator;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 
 /**
@@ -59,8 +60,8 @@ public abstract Object filterInferenceValue(
     );
 
     /**
-     * Abstract helper method to filter individual values based on the existing and new metadata maps.
-     * Implementations should provide logic to compare values and determine if embeddings can be reused.
+     * Abstract helper method to filter individual objects based on the existing and new metadata maps.
+     * Implementations should provide logic to compare objects and determine if embeddings can be reused.
      *
      * @param embeddingKey The dot-notation path for the embedding field
      * @param processValue The value to be checked for potential embedding reuse
@@ -71,7 +72,7 @@ public abstract Object filterInferenceValue(
      * @return The processed value or null if embeddings are reused
      */
 
-    public abstract Object copyEmbeddingForSingleValue(
+    public abstract Object copyEmbeddingForSingleObject(
         String embeddingKey,
         Object processValue,
         Object existingValue,
@@ -81,9 +82,9 @@ public abstract Object copyEmbeddingForSingleValue(
     );
 
     /**
-     * Abstract method to filter and compare lists of values.
-     * If all elements in the list are identical between the new and existing metadata maps, embeddings are copied,
-     * and an empty list is returned to indicate no further processing is required.
+     * Abstract method to filter and compare lists of objects.
+     * If all objects in the list are identical between the new and existing metadata maps, embeddings are copied,
+     * and null is returned to indicate no further processing is required.
      *
      * @param embeddingKey The dot-notation path for the embedding field
      * @param processList The list of values to be checked for potential embedding reuse
@@ -93,7 +94,7 @@ public abstract Object copyEmbeddingForSingleValue(
      * @return A processed list or an empty list if embeddings are reused.
      */
 
-    public abstract List<Object> copyEmbeddingForMultipleValues(
+    public abstract List<Object> copyEmbeddingForListObject(
         String embeddingKey,
         List<Object> processList,
         List<Object> existingList,
@@ -194,7 +195,7 @@ protected List<Object> filterListValue(
         List<Object> existingListValue = ProcessorUtils.unsafeCastToObjectList(existingListOptional.get());
         if (existingListValue.getFirst() instanceof List) {
             // in case of nested list, compare and copy by list comparison
-            return copyEmbeddingForMultipleValues(
+            return copyEmbeddingForListObject(
                 embeddingKey,
                 processList,
                 ProcessorUtils.unsafeCastToObjectList(existingListValue.getFirst()),
@@ -235,16 +236,19 @@ public List<Object> filterMapValuesInList(
         ListIterator<Object> existingListIterator = existingList.listIterator();
         ListIterator<Object> embeddingListIterator = embeddingList.listIterator();
         int index = 0;
-        while (processListIterator.hasNext() && existingListIterator.hasNext() && embeddingListIterator.hasNext()) {
-            Object processedItem = copyEmbeddingForSingleValue(
-                embeddingKey,
-                processListIterator.next(),
-                existingListIterator.next(),
-                embeddingListIterator.next(),
-                sourceAndMetadataMap,
-                index++
-            );
-            filteredList.add(processedItem);
+        for (Object processValue : processList) {
+            if (Objects.nonNull(processValue) && existingListIterator.hasNext() && embeddingListIterator.hasNext()) {
+                Object processedItem = copyEmbeddingForSingleObject(
+                    embeddingKey,
+                    processValue,
+                    existingListIterator.next(),
+                    embeddingListIterator.next(),
+                    sourceAndMetadataMap,
+                    index
+                );
+                filteredList.add(processedItem);
+            }
+            index++;
         }
         return filteredList;
     }

src/main/java/org/opensearch/neuralsearch/processor/optimization/InferenceFilter.java

@@ -7,7 +7,6 @@
 import lombok.extern.log4j.Log4j2;
 import org.opensearch.neuralsearch.processor.util.ProcessorUtils;
 
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -47,7 +46,7 @@ public Object filterInferenceValue(
         Optional<Object> existingValueOptional = ProcessorUtils.getValueFromSource(existingSourceAndMetadataMap, textPath);
         Optional<Object> embeddingValueOptional = ProcessorUtils.getValueFromSource(existingSourceAndMetadataMap, embeddingKey);
         if (existingValueOptional.isPresent() && embeddingValueOptional.isPresent()) {
-            return copyEmbeddingForSingleValue(
+            return copyEmbeddingForSingleObject(
                 embeddingKey,
                 processValue,
                 existingValueOptional.get(),
@@ -67,7 +66,7 @@ public Object filterInferenceValue(
      * @return null if embeddings are reused; the processValue otherwise.
      */
     @Override
-    public Object copyEmbeddingForSingleValue(
+    public Object copyEmbeddingForSingleObject(
         String embeddingKey,
         Object processValue,
         Object existingValue,
@@ -90,7 +89,7 @@ public Object copyEmbeddingForSingleValue(
      * @return empty list if embeddings are reused; processList otherwise.
      */
     @Override
-    public List<Object> copyEmbeddingForMultipleValues(
+    public List<Object> copyEmbeddingForListObject(
         String embeddingKey,
         List<Object> processList,
         List<Object> existingList,
@@ -99,8 +98,8 @@ public List<Object> copyEmbeddingForMultipleValues(
     ) {
         if (Objects.equals(processList, existingList)) {
             ProcessorUtils.setValueToSource(sourceAndMetadataMap, embeddingKey, embeddingList);
-            // if successfully copied, return empty list to be filtered out from process map
-            return Collections.emptyList();
+            // if successfully copied, return null to be filtered out from process map
+            return null;
         }
         // source list and existing list are different, return processList to be included in process map
         return processList;

src/main/java/org/opensearch/neuralsearch/processor/optimization/TextEmbeddingInferenceFilter.java

@@ -306,12 +306,12 @@ public static boolean isNumeric(Object value) {
      */
     public static int getNumOfSubqueries(final List<CompoundTopDocs> queryTopDocs) {
         return queryTopDocs.stream()
-                .filter(Objects::nonNull)
-                .filter(topDocs -> !topDocs.getTopDocs().isEmpty())
-                .findAny()
-                .get()
-                .getTopDocs()
-                .size();
+            .filter(Objects::nonNull)
+            .filter(topDocs -> !topDocs.getTopDocs().isEmpty())
+            .findAny()
+            .get()
+            .getTopDocs()
+            .size();
     }
 
     // This method should be used only when you are certain the object is a `Map<String, Object>`.

src/main/java/org/opensearch/neuralsearch/processor/util/ProcessorUtils.java

@@ -7,6 +7,8 @@
 import com.google.common.collect.ImmutableList;
 import org.apache.commons.lang.math.RandomUtils;
 import org.opensearch.action.get.GetResponse;
+import org.opensearch.action.get.MultiGetItemResponse;
+import org.opensearch.action.get.MultiGetResponse;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.XContentBuilder;
@@ -24,18 +26,22 @@
 import java.util.Map;
 
 public class InferenceProcessorTestCase extends OpenSearchTestCase {
-
-    protected List<IngestDocumentWrapper> createIngestDocumentWrappers(int count) {
+    protected List<IngestDocumentWrapper> createIngestDocumentWrappers(int count, String value) {
         List<IngestDocumentWrapper> wrapperList = new ArrayList<>();
-        for (int i = 0; i < count; ++i) {
+        for (int i = 1; i <= count; ++i) {
             Map<String, Object> sourceAndMetadata = new HashMap<>();
-            sourceAndMetadata.put("key1", "value1");
+            sourceAndMetadata.put("key1", value);
             sourceAndMetadata.put(IndexFieldMapper.NAME, "my_index");
+            sourceAndMetadata.put("_id", String.valueOf(i));
             wrapperList.add(new IngestDocumentWrapper(i, new IngestDocument(sourceAndMetadata, new HashMap<>()), null));
         }
         return wrapperList;
     }
 
+    protected List<IngestDocumentWrapper> createIngestDocumentWrappers(int count) {
+        return createIngestDocumentWrappers(count, "value1");
+    }
+
     protected List<List<Float>> createMockVectorWithLength(int size) {
         float suffix = .234f;
         List<List<Float>> result = new ArrayList<>();
@@ -120,6 +126,10 @@ protected GetResponse mockEmptyGetResponse() throws IOException {
         return GetResponse.fromXContent(contentParser);
     }
 
+    protected MultiGetResponse mockEmptyMultiGetItemResponse() throws IOException {
+        return new MultiGetResponse(new MultiGetItemResponse[0]);
+    }
+
     protected GetResponse convertToGetResponse(IngestDocument ingestDocument) throws IOException {
         String index = ingestDocument.getSourceAndMetadata().get("_index").toString();
         String id = ingestDocument.getSourceAndMetadata().get("_id").toString();
@@ -130,4 +140,12 @@ protected GetResponse convertToGetResponse(IngestDocument ingestDocument) throws
         GetResult result = new GetResult(index, id, 0, 1, 1, true, bytes, null, null);
         return new GetResponse(result);
     }
+
+    protected MultiGetResponse convertToMultiGetItemResponse(List<IngestDocumentWrapper> ingestDocuments) throws IOException {
+        MultiGetItemResponse[] multiGetItems = new MultiGetItemResponse[ingestDocuments.size()];
+        for (int i = 0; i < ingestDocuments.size(); i++) {
+            multiGetItems[i] = new MultiGetItemResponse(convertToGetResponse(ingestDocuments.get(i).getIngestDocument()), null);
+        }
+        return new MultiGetResponse(multiGetItems);
+    }
 }