Adding non empty check before filling in result

Signed-off-by: wangdongyu.danny <[email protected]>
opensearch-project · Sep 20, 2024 · 22e99ab · 22e99ab
1 parent 481a347
commit 22e99ab
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Enhancements
 - Set neural-search plugin 3.0.0 baseline JDK version to JDK-2 ([#838](https://github.com/opensearch-project/neural-search/pull/838))
 ### Bug Fixes
+- Fix for nested field missing sub embedding field in text embedding processor
 ### Infrastructure
 ### Documentation
 ### Maintenance

diff --git a/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/InferenceProcessor.java
@@ -419,8 +419,14 @@ private void putNLPResultToSourceMapForMapType(
             for (Map.Entry<String, Object> inputNestedMapEntry : ((Map<String, Object>) sourceValue).entrySet()) {
                 if (sourceAndMetadataMap.get(processorKey) instanceof List) {
                     // build nlp output for list of nested objects
+                    List<Object> inputNestedMapValue = (List<Object>) inputNestedMapEntry.getValue();
+                    int valueIndex = 0;
                     for (Map<String, Object> nestedElement : (List<Map<String, Object>>) sourceAndMetadataMap.get(processorKey)) {
-                        nestedElement.put(inputNestedMapEntry.getKey(), results.get(indexWrapper.index++));
+                        // Only fill in when value is not null
+                        if (valueIndex < inputNestedMapValue.size() && inputNestedMapValue.get(valueIndex) != null) {
+                            nestedElement.put(inputNestedMapEntry.getKey(), results.get(indexWrapper.index++));
+                        }
+                        valueIndex++;
                     }
                 } else {
                     Pair<String, Object> processedNestedKey = processNestedKey(inputNestedMapEntry);

diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java
@@ -739,6 +739,20 @@ public void testBuildVectorOutput_withNestedList_successful() {
         assertNotNull(nestedObj.get(1).get("vectorField"));
     }
 
+    public void testBuildVectorOutput_withNestedListHasNotForEmbeddingField_successful() {
+        Map<String, Object> config = createNestedListConfiguration();
+        IngestDocument ingestDocument = createNestedListWithNotEmbeddingFieldIngestDocument();
+        TextEmbeddingProcessor textEmbeddingProcessor = createInstanceWithNestedMapConfiguration(config);
+        Map<String, Object> knnMap = textEmbeddingProcessor.buildMapWithTargetKeys(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");
+        assertFalse(nestedObj.get(0).containsKey("vectorField"));
+        assertTrue(nestedObj.get(0).containsKey("textFieldNotForEmbedding"));
+        assertTrue(nestedObj.get(1).containsKey("vectorField"));
+        assertNotNull(nestedObj.get(1).get("vectorField"));
+    }
+
     public void testBuildVectorOutput_withNestedList_Level2_successful() {
         Map<String, Object> config = createNestedList2LevelConfiguration();
         IngestDocument ingestDocument = create2LevelNestedListIngestDocument();
@@ -754,6 +768,21 @@ public void testBuildVectorOutput_withNestedList_Level2_successful() {
         assertNotNull(nestedObj.get(1).get("vectorField"));
     }
 
+    public void testBuildVectorOutput_withNestedListHasNotForEmbeddingField_Level2_successful() {
+        Map<String, Object> config = createNestedList2LevelConfiguration();
+        IngestDocument ingestDocument = create2LevelNestedListWithNotEmbeddingFieldIngestDocument();
+        TextEmbeddingProcessor textEmbeddingProcessor = createInstanceWithNestedMapConfiguration(config);
+        Map<String, Object> knnMap = textEmbeddingProcessor.buildMapWithTargetKeys(ingestDocument);
+        List<List<Float>> modelTensorList = createMockVectorResult();
+        textEmbeddingProcessor.buildNLPResult(knnMap, modelTensorList, ingestDocument.getSourceAndMetadata());
+        Map<String, Object> nestedLevel1 = (Map<String, Object>) ingestDocument.getSourceAndMetadata().get("nestedField");
+        List<Map<String, Object>> nestedObj = (List<Map<String, Object>>) nestedLevel1.get("nestedField");
+        assertFalse(nestedObj.get(0).containsKey("vectorField"));
+        assertTrue(nestedObj.get(0).containsKey("textFieldNotForEmbedding"));
+        assertTrue(nestedObj.get(1).containsKey("vectorField"));
+        assertNotNull(nestedObj.get(1).get("vectorField"));
+    }
+
     public void test_updateDocument_appendVectorFieldsToDocument_successful() {
         Map<String, Object> config = createPlainStringConfiguration();
         IngestDocument ingestDocument = createPlainIngestDocument();
@@ -1039,6 +1068,16 @@ private IngestDocument createNestedListIngestDocument() {
         return new IngestDocument(nestedList, new HashMap<>());
     }
 
+    private IngestDocument createNestedListWithNotEmbeddingFieldIngestDocument() {
+        HashMap<String, Object> nestedObj1 = new HashMap<>();
+        nestedObj1.put("textFieldNotForEmbedding", "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<>());
+    }
+
     private IngestDocument create2LevelNestedListIngestDocument() {
         HashMap<String, Object> nestedObj1 = new HashMap<>();
         nestedObj1.put("textField", "This is a text field");
@@ -1050,4 +1089,16 @@ private IngestDocument create2LevelNestedListIngestDocument() {
         nestedList1.put("nestedField", nestedList);
         return new IngestDocument(nestedList1, new HashMap<>());
     }
+
+    private IngestDocument create2LevelNestedListWithNotEmbeddingFieldIngestDocument() {
+        HashMap<String, Object> nestedObj1 = new HashMap<>();
+        nestedObj1.put("textFieldNotForEmbedding", "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));
+        HashMap<String, Object> nestedList1 = new HashMap<>();
+        nestedList1.put("nestedField", nestedList);
+        return new IngestDocument(nestedList1, new HashMap<>());
+    }
 }
diff --git a/src/test/resources/processor/IndexMappings.json b/src/test/resources/processor/IndexMappings.json
@@ -90,6 +90,9 @@
           "text": {
             "type": "text"
           },
+          "text_not_for_embedding": {
+            "type": "text"
+          },
           "embedding": {
             "type": "knn_vector",
             "dimension": 768,

diff --git a/src/test/resources/processor/ingest_doc1.json b/src/test/resources/processor/ingest_doc1.json
@@ -12,6 +12,9 @@
     "movie": null
   },
   "nested_passages": [
+    {
+      "text_not_for_embedding": "test"
+    },
     {
       "text": "hello"
     },

diff --git a/src/test/resources/processor/ingest_doc2.json b/src/test/resources/processor/ingest_doc2.json
@@ -10,6 +10,9 @@
     "movie": null
   },
   "nested_passages": [
+    {
+      "text_not_for_embedding": "test"
+    },
     {
       "text": "apple"
     },