implement batch update document optimization in text embedding processor

Signed-off-by: will-hwang <[email protected]>

Author: will-hwang

CHANGELOG.md

@@ -5,7 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 3.0](https://github.com/opensearch-project/neural-search/compare/2.x...HEAD)
 ### Features
-- Add Optimized Text Embedding Processor ([#1191](https://github.com/opensearch-project/neural-search/pull/1191))
+- Optimizing embedding generation in text embedding processor ([#1191](https://github.com/opensearch-project/neural-search/pull/1191))
 ### Enhancements
 - Set neural-search plugin 3.0.0 baseline JDK version to JDK-21 ([#838](https://github.com/opensearch-project/neural-search/pull/838))
 - Support different embedding types in model's response ([#1007](https://github.com/opensearch-project/neural-search/pull/1007))

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

@@ -26,6 +26,8 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
+import org.opensearch.action.get.MultiGetItemResponse;
+import org.opensearch.action.get.MultiGetRequest;
 import org.opensearch.common.collect.Tuple;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.util.CollectionUtils;
@@ -54,6 +56,8 @@ public abstract class InferenceProcessor extends AbstractBatchingProcessor {
 
     public static final String MODEL_ID_FIELD = "model_id";
     public static final String FIELD_MAP_FIELD = "field_map";
+    public static final String INDEX_FIELD = "_index";
+    public static final String ID_FIELD = "_id";
     private static final BiFunction<Object, Object, Object> REMAPPING_FUNCTION = (v1, v2) -> {
         if (v1 instanceof Collection && v2 instanceof Collection) {
             ((Collection) v1).addAll((Collection) v2);
@@ -182,6 +186,15 @@ public void subBatchExecute(List<IngestDocumentWrapper> ingestDocumentWrappers,
             handler.accept(ingestDocumentWrappers);
             return;
         }
+        doSubBatchExecute(ingestDocumentWrappers, inferenceList, dataForInferences, handler);
+    }
+
+    protected void doSubBatchExecute(
+        List<IngestDocumentWrapper> ingestDocumentWrappers,
+        List<String> inferenceList,
+        List<DataForInference> dataForInferences,
+        Consumer<List<IngestDocumentWrapper>> handler
+    ) {
         Tuple<List<String>, Map<Integer, Integer>> sortedResult = sortByLengthAndReturnOriginalOrder(inferenceList);
         inferenceList = sortedResult.v1();
         Map<Integer, Integer> originalOrder = sortedResult.v2();
@@ -238,7 +251,7 @@ private List<?> restoreToOriginalOrder(List<?> results, Map<Integer, Integer> or
         return sortedResults;
     }
 
-    private List<String> constructInferenceTexts(List<DataForInference> dataForInferences) {
+    protected List<String> constructInferenceTexts(List<DataForInference> dataForInferences) {
         List<String> inferenceTexts = new ArrayList<>();
         for (DataForInference dataForInference : dataForInferences) {
             if (dataForInference.getIngestDocumentWrapper().getException() != null
@@ -250,7 +263,7 @@ private List<String> constructInferenceTexts(List<DataForInference> dataForInfer
         return inferenceTexts;
     }
 
-    private List<DataForInference> getDataForInference(List<IngestDocumentWrapper> ingestDocumentWrappers) {
+    protected List<DataForInference> getDataForInference(List<IngestDocumentWrapper> ingestDocumentWrappers) {
         List<DataForInference> dataForInferences = new ArrayList<>();
         for (IngestDocumentWrapper ingestDocumentWrapper : ingestDocumentWrappers) {
             Map<String, Object> processMap = null;
@@ -272,7 +285,7 @@ private List<DataForInference> getDataForInference(List<IngestDocumentWrapper> i
 
     @Getter
     @AllArgsConstructor
-    private static class DataForInference {
+    protected static class DataForInference {
         private final IngestDocumentWrapper ingestDocumentWrapper;
         private final Map<String, Object> processMap;
         private final List<String> inferenceList;
@@ -415,6 +428,36 @@ protected void setVectorFieldsToDocument(IngestDocument ingestDocument, Map<Stri
         nlpResult.forEach(ingestDocument::setFieldValue);
     }
 
+    /**
+     * This method creates a MultiGetRequest from a list of ingest documents to be fetched for comparison
+     * @param ingestDocumentWrappers, list of ingest documents
+     * */
+    protected MultiGetRequest buildMultiGetRequest(List<IngestDocumentWrapper> ingestDocumentWrappers) {
+        MultiGetRequest multiGetRequest = new MultiGetRequest();
+        for (IngestDocumentWrapper ingestDocumentWrapper : ingestDocumentWrappers) {
+            Object index = ingestDocumentWrapper.getIngestDocument().getSourceAndMetadata().get(INDEX_FIELD);
+            Object id = ingestDocumentWrapper.getIngestDocument().getSourceAndMetadata().get(ID_FIELD);
+            if (Objects.nonNull(index) && Objects.nonNull(id)) {
+                multiGetRequest.add(index.toString(), id.toString());
+            }
+        }
+        return multiGetRequest;
+    }
+
+    /**
+     * This method creates a map of documents from MultiGetItemResponse where the key is document ID and value is corresponding document
+     * @param multiGetItemResponses, array of responses from Multi Get Request
+     * */
+    protected Map<String, Map<String, Object>> createDocumentMap(MultiGetItemResponse[] multiGetItemResponses) {
+        Map<String, Map<String, Object>> existingDocuments = new HashMap<>();
+        for (MultiGetItemResponse item : multiGetItemResponses) {
+            String id = item.getId();
+            Map<String, Object> existingDocument = item.getResponse().getSourceAsMap();
+            existingDocuments.put(id, existingDocument);
+        }
+        return existingDocuments;
+    }
+
     @SuppressWarnings({ "unchecked" })
     @VisibleForTesting
     Map<String, Object> buildNLPResult(Map<String, Object> processorMap, List<?> results, Map<String, Object> sourceAndMetadataMap) {
@@ -582,7 +625,7 @@ private List<Map<String, Object>> buildNLPResultForListType(List<String> sourceV
         List<Map<String, Object>> keyToResult = new ArrayList<>();
         sourceValue.stream()
             .filter(Objects::nonNull) // explicit null check is required since sourceValue can contain null values in cases where
-                                      // sourceValue has been filtered
+            // sourceValue has been filtered
             .forEachOrdered(x -> keyToResult.add(ImmutableMap.of(listTypeNestedMapKey, results.get(indexWrapper.index++))));
         return keyToResult;
     }

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

@@ -4,6 +4,8 @@
  */
 package org.opensearch.neuralsearch.processor;
 
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -13,10 +15,14 @@
 
 import org.opensearch.action.get.GetAction;
 import org.opensearch.action.get.GetRequest;
+import org.opensearch.action.get.MultiGetAction;
+import org.opensearch.action.get.MultiGetItemResponse;
 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;
@@ -106,4 +112,55 @@ 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) {
+        if (CollectionUtils.isEmpty(ingestDocumentWrappers)) {
+            handler.accept(Collections.emptyList());
+            return;
+        }
+
+        List<DataForInference> dataForInferences = getDataForInference(ingestDocumentWrappers);
+        List<String> inferenceList = constructInferenceTexts(dataForInferences);
+        if (inferenceList.isEmpty()) {
+            handler.accept(ingestDocumentWrappers);
+            return;
+        }
+        if (skipExisting == false) {
+            doSubBatchExecute(ingestDocumentWrappers, inferenceList, dataForInferences, handler);
+            return;
+        }
+        openSearchClient.execute(MultiGetAction.INSTANCE, buildMultiGetRequest(ingestDocumentWrappers), ActionListener.wrap(response -> {
+            MultiGetItemResponse[] multiGetItemResponses = response.getResponses();
+            if (multiGetItemResponses == null || multiGetItemResponses.length == 0) {
+                doSubBatchExecute(ingestDocumentWrappers, inferenceList, dataForInferences, handler);
+                return;
+            }
+            Map<String, Map<String, Object>> existingDocuments = createDocumentMap(multiGetItemResponses);
+            List<DataForInference> filteredDataForInference = new ArrayList<>();
+            for (DataForInference dataForInference : dataForInferences) {
+                IngestDocumentWrapper ingestDocumentWrapper = dataForInference.getIngestDocumentWrapper();
+                Map<String, Object> processMap = dataForInference.getProcessMap();
+                Map<String, Object> document = ingestDocumentWrapper.getIngestDocument().getSourceAndMetadata();
+                Object id = document.get(ID_FIELD);
+                // insert non-filtered dataForInference if existing document does not exist
+                if (Objects.isNull(id) || existingDocuments.containsKey(id.toString()) == false) {
+                    filteredDataForInference.add(dataForInference);
+                    continue;
+                }
+                // filter dataForInference when existing document exists
+                String docId = id.toString();
+                Map<String, Object> existingDocument = existingDocuments.get(docId);
+                Map<String, Object> filteredProcessMap = textEmbeddingInferenceFilter.filter(existingDocument, document, processMap);
+                List<String> filteredInferenceList = createInferenceList(filteredProcessMap);
+                filteredDataForInference.add(new DataForInference(ingestDocumentWrapper, filteredProcessMap, filteredInferenceList));
+            }
+            List<String> filteredInferenceList = constructInferenceTexts(filteredDataForInference);
+            if (filteredInferenceList.isEmpty()) {
+                handler.accept(ingestDocumentWrappers);
+            } else {
+                doSubBatchExecute(ingestDocumentWrappers, filteredInferenceList, filteredDataForInference, handler);
+            }
+        }, e -> { handler.accept(null); }));
+    }
 }

src/main/java/org/opensearch/neuralsearch/processor/optimization/InferenceFilter.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;
 
 /**
@@ -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 copyEmbedding(
         String embeddingKey,
         Object processValue,
         Object existingValue,
@@ -80,27 +81,6 @@ public abstract Object copyEmbeddingForSingleValue(
         int index
     );
 
-    /**
-     * 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.
-     *
-     * @param embeddingKey The dot-notation path for the embedding field
-     * @param processList The list of values to be checked for potential embedding reuse
-     * @param existingList The list of existing values for comparison
-     * @param embeddingList The list of existing embeddings
-     * @param sourceAndMetadataMap The metadata map of the new document.
-     * @return A processed list or an empty list if embeddings are reused.
-     */
-
-    public abstract List<Object> copyEmbeddingForMultipleValues(
-        String embeddingKey,
-        List<Object> processList,
-        List<Object> existingList,
-        List<Object> embeddingList,
-        Map<String, Object> sourceAndMetadataMap
-    );
-
     /**
      * This method navigates through the nested structure, checking each key-value pair recursively. It supports:
      * Map values: Processed recursively using this method.
@@ -155,7 +135,7 @@ private Map<String, Object> filter(
                 );
                 filteredProcessMap.put(key, filteredInnerMap.isEmpty() ? null : filteredInnerMap);
             } else if (value instanceof List) {
-                List<Object> processedList = filterListValue(
+                Object processedList = filterListValue(
                     currentPath,
                     ProcessorUtils.unsafeCastToObjectList(value),
                     sourceAndMetadataMap,
@@ -194,13 +174,15 @@ 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(
+            Object processedList = copyEmbedding(
                 embeddingKey,
                 processList,
-                ProcessorUtils.unsafeCastToObjectList(existingListValue.getFirst()),
-                ProcessorUtils.unsafeCastToObjectList(embeddingListOptional.get()),
-                sourceAndMetadataMap
+                existingListValue.getFirst(),
+                embeddingListOptional.get(),
+                sourceAndMetadataMap,
+                -1
             );
+            return Objects.nonNull(processedList) ? ProcessorUtils.unsafeCastToObjectList(processedList) : null;
         } else {
             // in case of List of Maps, compare each map entry in list
             return filterMapValuesInList(
@@ -231,20 +213,22 @@ public List<Object> filterMapValuesInList(
         Map<String, Object> sourceAndMetadataMap
     ) {
         List<Object> filteredList = new ArrayList<>();
-        ListIterator<Object> processListIterator = processList.listIterator();
         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 = copyEmbedding(
+                    embeddingKey,
+                    processValue,
+                    existingListIterator.next(),
+                    embeddingListIterator.next(),
+                    sourceAndMetadataMap,
+                    index
+                );
+                filteredList.add(processedItem);
+            }
+            index++;
         }
         return filteredList;
     }

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

@@ -7,8 +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;
 import java.util.Optional;
@@ -47,7 +45,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 copyEmbedding(
                 embeddingKey,
                 processValue,
                 existingValueOptional.get(),
@@ -60,14 +58,14 @@ public Object filterInferenceValue(
     }
 
     /**
-     * Copy a single value by checking if the text is identical in both the existing and new document.
+     * Copy a single value by checking if the given texts is identical in both the existing and new document.
      * If the text matches, the corresponding embedding is copied, and null is returned, indicating no further
      * processing is required.
      *
      * @return null if embeddings are reused; the processValue otherwise.
      */
     @Override
-    public Object copyEmbeddingForSingleValue(
+    public Object copyEmbedding(
         String embeddingKey,
         Object processValue,
         Object existingValue,
@@ -83,26 +81,4 @@ public Object copyEmbeddingForSingleValue(
         // processValue and existingValue are different, return processValue to be included in process map
         return processValue;
     }
-
-    /**
-     * Copy values in list by checking if all texts in list are identical in both the existing and new documents.
-     * If lists are equal, the corresponding embeddings are copied
-     * @return empty list if embeddings are reused; processList otherwise.
-     */
-    @Override
-    public List<Object> copyEmbeddingForMultipleValues(
-        String embeddingKey,
-        List<Object> processList,
-        List<Object> existingList,
-        List<Object> embeddingList,
-        Map<String, Object> sourceAndMetadataMap
-    ) {
-        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();
-        }
-        // source list and existing list are different, return processList to be included in process map
-        return processList;
-    }
 }