implement batch document update scenario for text embedding processor

Signed-off-by: Will Hwang <sang7239@gmail.com>

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))

DEVELOPER_GUIDE.md

@@ -351,9 +351,9 @@ through the same build issue.
 
 ### Class and package names
 
-Class names should use `CamelCase`. 
+Class names should use `CamelCase`.
 
-Try to put new classes into existing packages if package name abstracts the purpose of the class. 
+Try to put new classes into existing packages if package name abstracts the purpose of the class.
 
 Example of good class file name and package utilization:
 
@@ -371,7 +371,7 @@ methods rather than a long single one and does everything.
 ### Documentation
 
 Document you code. That includes purpose of new classes, every public method and code sections that have critical or non-trivial
-logic (check this example https://github.com/opensearch-project/neural-search/blob/main/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java#L238). 
+logic (check this example https://github.com/opensearch-project/neural-search/blob/main/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java#L238).
 
 When you submit a feature PR, please submit a new
 [documentation issue](https://github.com/opensearch-project/documentation-website/issues/new/choose). This is a path for the documentation to be published as part of https://opensearch.org/docs/latest/ documentation site.
@@ -384,17 +384,17 @@ For the most part, we're using common conventions for Java projects. Here are a
 
 1. Use descriptive names for classes, methods, fields, and variables.
 2. Avoid abbreviations unless they are widely accepted
-3. Use `final` on all method arguments unless it's absolutely necessary 
+3. Use `final` on all method arguments unless it's absolutely necessary
 4. Wildcard imports are not allowed.
 5. Static imports are preferred over qualified imports when using static methods
 6. Prefer creating non-static public methods whenever possible. Avoid static methods in general, as they can often serve as shortcuts.
 Static methods are acceptable if they are private and do not access class state.
-7. Use functional programming style inside methods unless it's a performance critical section. 
+7. Use functional programming style inside methods unless it's a performance critical section.
 8. For parameters of lambda expression please use meaningful names instead of shorten cryptic ones.
 9. Use Optional for return values if the value may not be present. This should be preferred to returning null.
 10. Do not create checked exceptions, and do not throw checked exceptions from public methods whenever possible. In general, if you call a method with a checked exception, you should wrap that exception into an unchecked exception.
 11. Throwing checked exceptions from private methods is acceptable.
-12. Use String.format when a string includes parameters, and prefer this over direct string concatenation. Always specify a Locale with String.format; 
+12. Use String.format when a string includes parameters, and prefer this over direct string concatenation. Always specify a Locale with String.format;
 as a rule of thumb, use Locale.ROOT.
 13. Prefer Lombok annotations to the manually written boilerplate code
 14. When throwing an exception, avoid including user-provided content in the exception message. For secure coding practices,
@@ -440,17 +440,17 @@ Fix any new warnings before submitting your PR to ensure proper code documentati
 
 ### Tests
 
-Write unit and integration tests for your new functionality. 
+Write unit and integration tests for your new functionality.
 
 Unit tests are preferred as they are cheap and fast, try to use them to cover all possible
-combinations of parameters. Utilize mocks to mimic dependencies. 
+combinations of parameters. Utilize mocks to mimic dependencies.
 
-Integration tests should be used sparingly, focusing primarily on the main (happy path) scenario or cases where extensive 
-mocking is impractical. Include one or two unhappy paths to confirm that correct response codes are returned to the user. 
-Whenever possible, favor scenarios that do not require model deployment. If model deployment is necessary, use an existing 
+Integration tests should be used sparingly, focusing primarily on the main (happy path) scenario or cases where extensive
+mocking is impractical. Include one or two unhappy paths to confirm that correct response codes are returned to the user.
+Whenever possible, favor scenarios that do not require model deployment. If model deployment is necessary, use an existing
 model, as tests involving new model deployments are the most resource-intensive.
 
-If your changes could affect backward compatibility, please include relevant backward compatibility tests along with your 
+If your changes could affect backward compatibility, please include relevant backward compatibility tests along with your
 PR. For guidance on adding these tests, refer to the [Backwards Compatibility Testing](#backwards-compatibility-testing) section in this guide.
 
 ### Outdated or irrelevant code

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);
@@ -169,49 +173,80 @@ void preprocessIngestDocument(IngestDocument ingestDocument) {
      */
     abstract void doBatchExecute(List<String> inferenceList, Consumer<List<?>> handler, Consumer<Exception> onException);
 
+    /**
+     * This is the function which does actual inference work for subBatchExecute interface.
+     * @param ingestDocumentWrappers a list of IngestDocuments in a batch.
+     * @param handler a callback handler to handle inference results which is a list of objects.
+     */
     @Override
     public void subBatchExecute(List<IngestDocumentWrapper> ingestDocumentWrappers, Consumer<List<IngestDocumentWrapper>> handler) {
-        if (CollectionUtils.isEmpty(ingestDocumentWrappers)) {
-            handler.accept(Collections.emptyList());
-            return;
-        }
+        try {
+            if (CollectionUtils.isEmpty(ingestDocumentWrappers)) {
+                handler.accept(Collections.emptyList());
+                return;
+            }
 
-        List<DataForInference> dataForInferences = getDataForInference(ingestDocumentWrappers);
-        List<String> inferenceList = constructInferenceTexts(dataForInferences);
-        if (inferenceList.isEmpty()) {
+            List<DataForInference> dataForInferences = getDataForInference(ingestDocumentWrappers);
+            List<String> inferenceList = constructInferenceTexts(dataForInferences);
+            if (inferenceList.isEmpty()) {
+                handler.accept(ingestDocumentWrappers);
+                return;
+            }
+            doSubBatchExecute(ingestDocumentWrappers, inferenceList, dataForInferences, handler);
+        } catch (Exception e) {
+            updateWithExceptions(ingestDocumentWrappers, e);
             handler.accept(ingestDocumentWrappers);
-            return;
         }
-        Tuple<List<String>, Map<Integer, Integer>> sortedResult = sortByLengthAndReturnOriginalOrder(inferenceList);
-        inferenceList = sortedResult.v1();
-        Map<Integer, Integer> originalOrder = sortedResult.v2();
-        doBatchExecute(inferenceList, results -> {
-            int startIndex = 0;
-            results = restoreToOriginalOrder(results, originalOrder);
-            for (DataForInference dataForInference : dataForInferences) {
-                if (dataForInference.getIngestDocumentWrapper().getException() != null
-                    || CollectionUtils.isEmpty(dataForInference.getInferenceList())) {
-                    continue;
+    }
+
+    /**
+     * This is a helper function for subBatchExecute, which invokes doBatchExecute for given inference list.
+     * @param ingestDocumentWrappers a list of IngestDocuments in a batch.
+     * @param inferenceList a list of String for inference.
+     * @param dataForInferences a list of data for inference, which includes ingestDocumentWrapper, processMap, inferenceList.
+     * @param handler a callback handler to handle inference results which is a list of objects.
+     */
+    protected void doSubBatchExecute(
+        List<IngestDocumentWrapper> ingestDocumentWrappers,
+        List<String> inferenceList,
+        List<DataForInference> dataForInferences,
+        Consumer<List<IngestDocumentWrapper>> handler
+    ) {
+        try {
+            Tuple<List<String>, Map<Integer, Integer>> sortedResult = sortByLengthAndReturnOriginalOrder(inferenceList);
+            inferenceList = sortedResult.v1();
+            Map<Integer, Integer> originalOrder = sortedResult.v2();
+            doBatchExecute(inferenceList, results -> {
+                int startIndex = 0;
+                results = restoreToOriginalOrder(results, originalOrder);
+                for (DataForInference dataForInference : dataForInferences) {
+                    if (dataForInference.getIngestDocumentWrapper().getException() != null
+                        || CollectionUtils.isEmpty(dataForInference.getInferenceList())) {
+                        continue;
+                    }
+                    List<?> inferenceResults = results.subList(startIndex, startIndex + dataForInference.getInferenceList().size());
+                    startIndex += dataForInference.getInferenceList().size();
+                    setVectorFieldsToDocument(
+                        dataForInference.getIngestDocumentWrapper().getIngestDocument(),
+                        dataForInference.getProcessMap(),
+                        inferenceResults
+                    );
                 }
-                List<?> inferenceResults = results.subList(startIndex, startIndex + dataForInference.getInferenceList().size());
-                startIndex += dataForInference.getInferenceList().size();
-                setVectorFieldsToDocument(
-                    dataForInference.getIngestDocumentWrapper().getIngestDocument(),
-                    dataForInference.getProcessMap(),
-                    inferenceResults
-                );
-            }
-            handler.accept(ingestDocumentWrappers);
-        }, exception -> {
-            for (IngestDocumentWrapper ingestDocumentWrapper : ingestDocumentWrappers) {
-                // The IngestDocumentWrapper might already run into exception and not sent for inference. So here we only
-                // set exception to IngestDocumentWrapper which doesn't have exception before.
-                if (ingestDocumentWrapper.getException() == null) {
-                    ingestDocumentWrapper.update(ingestDocumentWrapper.getIngestDocument(), exception);
+                handler.accept(ingestDocumentWrappers);
+            }, exception -> {
+                for (IngestDocumentWrapper ingestDocumentWrapper : ingestDocumentWrappers) {
+                    // The IngestDocumentWrapper might already run into exception and not sent for inference. So here we only
+                    // set exception to IngestDocumentWrapper which doesn't have exception before.
+                    if (ingestDocumentWrapper.getException() == null) {
+                        ingestDocumentWrapper.update(ingestDocumentWrapper.getIngestDocument(), exception);
+                    }
                 }
-            }
+                handler.accept(ingestDocumentWrappers);
+            });
+        } catch (Exception e) {
+            updateWithExceptions(ingestDocumentWrappers, e);
             handler.accept(ingestDocumentWrappers);
-        });
+        }
     }
 
     private Tuple<List<String>, Map<Integer, Integer>> sortByLengthAndReturnOriginalOrder(List<String> inferenceList) {
@@ -238,7 +273,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 +285,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 +307,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 +450,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) {
@@ -504,6 +569,13 @@ private void processMapEntryValue(
         }
     }
 
+    // This method updates each ingestDocument with exceptions
+    protected void updateWithExceptions(List<IngestDocumentWrapper> ingestDocumentWrappers, Exception e) {
+        for (IngestDocumentWrapper ingestDocumentWrapper : ingestDocumentWrappers) {
+            ingestDocumentWrapper.update(ingestDocumentWrapper.getIngestDocument(), e);
+        }
+    }
+
     private void processMapEntryValue(
         List<?> results,
         IndexWrapper indexWrapper,
@@ -582,7 +654,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;
     }