Add unit tests + small fixes

krishy91 · Jan 16, 2024 · 6e54ec6 · 6e54ec6
1 parent 861d3c0
commit 6e54ec6
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java
@@ -4,7 +4,11 @@
  */
 package org.opensearch.neuralsearch.processor;
 
-import java.util.*;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
 import java.util.function.BiConsumer;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -183,7 +187,7 @@ private void buildMapWithProcessorKeyAndOriginalValueForMapType(
                 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 HashMap();
+                    Map<String, Object> map = new LinkedHashMap<>();
                     map.put(nestedFieldMapEntry.getKey(), listOfStrings);
                     buildMapWithProcessorKeyAndOriginalValueForMapType(
                         nestedFieldMapEntry.getKey(),

diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java b/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<>());
+    }
 }