Handle case with nested list of objects

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

@@ -11,6 +11,7 @@
 import java.util.Objects;
 import java.util.function.BiConsumer;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 import org.apache.commons.lang3.StringUtils;
@@ -51,14 +52,14 @@
     private final Environment environment;
 
     public InferenceProcessor(
-        String tag,
-        String description,
-        String type,
-        String listTypeNestedMapKey,
-        String modelId,
-        Map<String, Object> fieldMap,
-        MLCommonsClientAccessor clientAccessor,
-        Environment environment
+            String tag,
+            String description,
+            String type,
+            String listTypeNestedMapKey,
+            String modelId,
+            Map<String, Object> fieldMap,
+            MLCommonsClientAccessor clientAccessor,
+            Environment environment
     ) {
         super(tag, description);
         this.type = type;
@@ -74,21 +75,21 @@
 
     private void validateEmbeddingConfiguration(Map<String, Object> fieldMap) {
         if (fieldMap == null
-            || fieldMap.size() == 0
-            || fieldMap.entrySet()
+                || fieldMap.size() == 0
+                || fieldMap.entrySet()
                 .stream()
                 .anyMatch(
-                    x -> StringUtils.isBlank(x.getKey()) || Objects.isNull(x.getValue()) || StringUtils.isBlank(x.getValue().toString())
+                        x -> StringUtils.isBlank(x.getKey()) || Objects.isNull(x.getValue()) || StringUtils.isBlank(x.getValue().toString())
                 )) {
             throw new IllegalArgumentException("Unable to create the processor as field_map has invalid key or value");
         }
     }
 
     public abstract void doExecute(
-        IngestDocument ingestDocument,
-        Map<String, Object> ProcessMap,
-        List<String> inferenceList,
-        BiConsumer<IngestDocument, Exception> handler
+            IngestDocument ingestDocument,
+            Map<String, Object> ProcessMap,
+            List<String> inferenceList,
+            BiConsumer<IngestDocument, Exception> handler
     );
 
     @Override
@@ -165,21 +166,36 @@
     }
 
     private void buildMapWithProcessorKeyAndOriginalValueForMapType(
-        String parentKey,
-        Object processorKey,
-        Map<String, Object> sourceAndMetadataMap,
-        Map<String, Object> treeRes
+            String parentKey,
+            Object processorKey,
+            Map<String, Object> sourceAndMetadataMap,
+            Map<String, Object> treeRes
     ) {
         if (processorKey == null || sourceAndMetadataMap == null) return;
         if (processorKey instanceof Map) {
             Map<String, Object> next = new LinkedHashMap<>();
-            for (Map.Entry<String, Object> nestedFieldMapEntry : ((Map<String, Object>) processorKey).entrySet()) {
-                buildMapWithProcessorKeyAndOriginalValueForMapType(
-                    nestedFieldMapEntry.getKey(),
-                    nestedFieldMapEntry.getValue(),
-                    (Map<String, Object>) sourceAndMetadataMap.get(parentKey),
-                    next
-                );
+            if (sourceAndMetadataMap.get(parentKey) instanceof Map) {
+                for (Map.Entry<String, Object> nestedFieldMapEntry : ((Map<String, Object>) processorKey).entrySet()) {
+                    buildMapWithProcessorKeyAndOriginalValueForMapType(
+                            nestedFieldMapEntry.getKey(),
+                            nestedFieldMapEntry.getValue(),
+                            (Map<String, Object>) sourceAndMetadataMap.get(parentKey),
+                            next
+                    );
+                }
+            } else if (sourceAndMetadataMap.get(parentKey) instanceof List) {
+                for (Map.Entry<String, Object> nestedFieldMapEntry : ((Map<String, Object>) processorKey).entrySet()) {
+                    List<Map<String, Object>> list = (List<Map<String, Object>>) sourceAndMetadataMap.get(parentKey);
+                    List<Object> listOfStrings = list.stream().map(x -> x.get(nestedFieldMapEntry.getKey())).collect(Collectors.toList());
+                    Map<String, Object> map = new LinkedHashMap<>();
+                    map.put(nestedFieldMapEntry.getKey(), listOfStrings);
+                    buildMapWithProcessorKeyAndOriginalValueForMapType(
+                            nestedFieldMapEntry.getKey(),
+                            nestedFieldMapEntry.getValue(),
+                            map,
+                            next
+                    );
+                }
             }
             treeRes.put(parentKey, next);
         } else {
@@ -196,7 +212,14 @@
                 String sourceKey = embeddingFieldsEntry.getKey();
                 Class<?> sourceValueClass = sourceValue.getClass();
                 if (List.class.isAssignableFrom(sourceValueClass) || Map.class.isAssignableFrom(sourceValueClass)) {
-                    validateNestedTypeValue(sourceKey, sourceValue, () -> 1);
+                    if (Map.class.isAssignableFrom(embeddingFieldsEntry.getValue().getClass())) {
+                        Map<String, Object> innerFieldsEntry = (Map<String, Object>) embeddingFieldsEntry.getValue();
+                        for (Map.Entry<String, Object> innerKey : innerFieldsEntry.entrySet()) {
+                            validateNestedTypeValue(innerKey.getKey(), sourceValue, () -> 2);
+                        }
+                    } else {
+                        validateNestedTypeValue(sourceKey, sourceValue, () -> 1);
+                    }
                 } else if (!String.class.isAssignableFrom(sourceValueClass)) {
                     throw new IllegalArgumentException("field [" + sourceKey + "] is neither string nor nested type, cannot process it");
                 } else if (StringUtils.isBlank(sourceValue.toString())) {
@@ -212,12 +235,12 @@
         if (maxDepth > MapperService.INDEX_MAPPING_DEPTH_LIMIT_SETTING.get(environment.settings())) {
             throw new IllegalArgumentException("map type field [" + sourceKey + "] reached max depth limit, cannot process it");
         } else if ((List.class.isAssignableFrom(sourceValue.getClass()))) {
-            validateListTypeValue(sourceKey, sourceValue);
+            validateListTypeValue(sourceKey, sourceValue, maxDepthSupplier);
         } else if (Map.class.isAssignableFrom(sourceValue.getClass())) {
             ((Map) sourceValue).values()
-                .stream()
-                .filter(Objects::nonNull)
-                .forEach(x -> validateNestedTypeValue(sourceKey, x, () -> maxDepth + 1));
+                    .stream()
+                    .filter(Objects::nonNull)
+                    .forEach(x -> validateNestedTypeValue(sourceKey, x, () -> maxDepth + 1));
         } else if (!String.class.isAssignableFrom(sourceValue.getClass())) {
             throw new IllegalArgumentException("map type field [" + sourceKey + "] has non-string type, cannot process it");
         } else if (StringUtils.isBlank(sourceValue.toString())) {
@@ -226,9 +249,11 @@
     }
 
     @SuppressWarnings({ "rawtypes" })
-    private void validateListTypeValue(String sourceKey, Object sourceValue) {
+    private void validateListTypeValue(String sourceKey, Object sourceValue, Supplier<Integer> maxDepthSupplier) {
         for (Object value : (List) sourceValue) {
-            if (value == null) {
+            if (value instanceof Map) {
+                validateNestedTypeValue(sourceKey, value, () -> maxDepthSupplier.get() + 1);
+            } else if (value == null) {
                 throw new IllegalArgumentException("list type field [" + sourceKey + "] has null, cannot process it");
             } else if (!(value instanceof String)) {
                 throw new IllegalArgumentException("list type field [" + sourceKey + "] has non string value, cannot process it");
@@ -266,22 +291,29 @@
 
     @SuppressWarnings({ "unchecked" })
     private void putNLPResultToSourceMapForMapType(
-        String processorKey,
-        Object sourceValue,
-        List<?> results,
-        IndexWrapper indexWrapper,
-        Map<String, Object> sourceAndMetadataMap
+            String processorKey,
+            Object sourceValue,
+            List<?> results,
+            IndexWrapper indexWrapper,
+            Map<String, Object> sourceAndMetadataMap
     ) {
         if (processorKey == null || sourceAndMetadataMap == null || sourceValue == null) return;
         if (sourceValue instanceof Map) {
             for (Map.Entry<String, Object> inputNestedMapEntry : ((Map<String, Object>) sourceValue).entrySet()) {
-                putNLPResultToSourceMapForMapType(
-                    inputNestedMapEntry.getKey(),
-                    inputNestedMapEntry.getValue(),
-                    results,
-                    indexWrapper,
-                    (Map<String, Object>) sourceAndMetadataMap.get(processorKey)
-                );
+                if (sourceAndMetadataMap.get(processorKey) instanceof List) {
+                    // build nlp output for list of nested objects
+                    for (Map<String, Object> nestedElement : (List<Map<String, Object>>) sourceAndMetadataMap.get(processorKey)) {
+                        nestedElement.put(inputNestedMapEntry.getKey(), results.get(indexWrapper.index++));
+                    }
+                } else {
+                    putNLPResultToSourceMapForMapType(
+                            inputNestedMapEntry.getKey(),
+                            inputNestedMapEntry.getValue(),
+                            results,
+                            indexWrapper,
+                            (Map<String, Object>) sourceAndMetadataMap.get(processorKey)
+                    );
+                }
             }
         } else if (sourceValue instanceof String) {
             sourceAndMetadataMap.put(processorKey, results.get(indexWrapper.index++));
@@ -293,7 +325,7 @@
     private List<Map<String, Object>> buildNLPResultForListType(List<String> sourceValue, List<?> results, IndexWrapper indexWrapper) {
         List<Map<String, Object>> keyToResult = new ArrayList<>();
         IntStream.range(0, sourceValue.size())
-            .forEachOrdered(x -> keyToResult.add(ImmutableMap.of(listTypeNestedMapKey, results.get(indexWrapper.index++))));
+                .forEachOrdered(x -> keyToResult.add(ImmutableMap.of(listTypeNestedMapKey, results.get(indexWrapper.index++))));
         return keyToResult;
     }
 

src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java

@@ -20,6 +20,7 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Arrays;
 import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 
@@ -404,6 +405,20 @@ public void testBuildVectorOutput_withNestedMap_successful() {
         assertNotNull(actionGamesKnn);
     }
 
+    public void testBuildVectorOutput_withNestedList_successful() {
+        Map<String, Object> config = createNestedListConfiguration();
+        IngestDocument ingestDocument = createNestedListIngestDocument();
+        TextEmbeddingProcessor textEmbeddingProcessor = createInstanceWithNestedMapConfiguration(config);
+        Map<String, Object> knnMap = textEmbeddingProcessor.buildMapWithProcessorKeyAndOriginalValue(ingestDocument);
+        List<List<Float>> modelTensorList = createMockVectorResult();
+        textEmbeddingProcessor.buildNLPResult(knnMap, modelTensorList, ingestDocument.getSourceAndMetadata());
+        List<Map<String, Object>> nestedObj = (List<Map<String, Object>>) ingestDocument.getSourceAndMetadata().get("nestedField");
+        assertTrue(nestedObj.get(0).containsKey("vectorField"));
+        assertTrue(nestedObj.get(1).containsKey("vectorField"));
+        assertNotNull(nestedObj.get(0).get("vectorField"));
+        assertNotNull(nestedObj.get(1).get("vectorField"));
+    }
+
     public void test_updateDocument_appendVectorFieldsToDocument_successful() {
         Map<String, Object> config = createPlainStringConfiguration();
         IngestDocument ingestDocument = createPlainIngestDocument();
@@ -520,4 +535,22 @@ private IngestDocument createNestedMapIngestDocument() {
         result.put("favorites", favorite);
         return new IngestDocument(result, new HashMap<>());
     }
+
+    private Map<String, Object> createNestedListConfiguration() {
+        Map<String, Object> nestedConfig = new HashMap<>();
+        nestedConfig.put("textField", "vectorField");
+        Map<String, Object> result = new HashMap<>();
+        result.put("nestedField", nestedConfig);
+        return result;
+    }
+
+    private IngestDocument createNestedListIngestDocument() {
+        HashMap<String, Object> nestedObj1 = new HashMap<>();
+        nestedObj1.put("textField", "This is a text field");
+        HashMap<String, Object> nestedObj2 = new HashMap<>();
+        nestedObj2.put("textField", "This is another text field");
+        HashMap<String, Object> nestedList = new HashMap<>();
+        nestedList.put("nestedField", Arrays.asList(nestedObj1, nestedObj2));
+        return new IngestDocument(nestedList, new HashMap<>());
+    }
 }