From 9264760e45d162895ff1f5c3970712572f152bdd Mon Sep 17 00:00:00 2001
From: Will Hwang <willhwang@Wills-Laptop.local>
Date: Tue, 18 Feb 2025 22:01:01 -0800
Subject: [PATCH] fix indexing in inference processor

---
 .../processor/InferenceProcessor.java         | 20 ++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java
index 6150e56eb..27d99a80c 100644
--- a/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java
+++ b/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java
@@ -560,7 +560,8 @@ private void processMapEntryValue(
     ) {
         // build nlp output for object in sourceValue which is map type
         Iterator<Map<String, Object>> iterator = sourceAndMetadataMapValueInList.iterator();
-        IntStream.range(0, sourceAndMetadataMapValueInList.size()).forEach(index -> {
+        IndexWrapper listIndexWrapper = new IndexWrapper(0);
+        for (int i = 0; i < sourceAndMetadataMapValueInList.size(); i++) {
             Map<String, Object> nestedElement = iterator.next();
             putNLPResultToSingleSourceMapInList(
                 inputNestedMapEntryKey,
@@ -568,9 +569,9 @@ private void processMapEntryValue(
                 results,
                 indexWrapper,
                 nestedElement,
-                index
+                listIndexWrapper
             );
-        });
+        }
     }
 
     /**
@@ -590,7 +591,7 @@ private void putNLPResultToSingleSourceMapInList(
         List<?> results,
         IndexWrapper indexWrapper,
         Map<String, Object> sourceAndMetadataMap,
-        int nestedElementIndex
+        IndexWrapper listIndexWrapper
     ) {
         if (processorKey == null || sourceAndMetadataMap == null || sourceValue == null) return;
         if (sourceValue instanceof Map) {
@@ -603,12 +604,17 @@ private void putNLPResultToSingleSourceMapInList(
                     results,
                     indexWrapper,
                     sourceMap,
-                    nestedElementIndex
+                    listIndexWrapper
                 );
             }
         } else {
-            if (sourceValue instanceof List && ((List<Object>) sourceValue).get(nestedElementIndex) != null) {
-                sourceAndMetadataMap.merge(processorKey, results.get(indexWrapper.index++), REMAPPING_FUNCTION);
+            if (sourceValue instanceof List) {
+                if (sourceAndMetadataMap.containsKey(processorKey)) {
+                    return;
+                }
+                if (((List<Object>) sourceValue).get(listIndexWrapper.index++) != null) {
+                    sourceAndMetadataMap.merge(processorKey, results.get(indexWrapper.index++), REMAPPING_FUNCTION);
+                }
             }
         }
     }