diff --git a/src/main/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessor.java b/src/main/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessor.java index 736fecafc..482a8730c 100644 --- a/src/main/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessor.java +++ b/src/main/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessor.java @@ -105,7 +105,7 @@ public void execute(IngestDocument ingestDocument, BiConsumer { - appendVectorFieldsToDocument(ingestDocument, knnMap, vectors); + setVectorFieldsToDocument(ingestDocument, knnMap, vectors); handler.accept(ingestDocument, null); }, e -> { handler.accept(null, e); })); } @@ -115,11 +115,11 @@ public void execute(IngestDocument ingestDocument, BiConsumer knnMap, List> vectors) { + void setVectorFieldsToDocument(IngestDocument ingestDocument, Map knnMap, List> vectors) { Objects.requireNonNull(vectors, "embedding failed, inference returns null result!"); log.debug("Text embedding result fetched, starting build vector output!"); Map textEmbeddingResult = buildTextEmbeddingResult(knnMap, vectors, ingestDocument.getSourceAndMetadata()); - textEmbeddingResult.forEach(ingestDocument::appendFieldValue); + textEmbeddingResult.forEach(ingestDocument::setFieldValue); } @SuppressWarnings({ "unchecked" }) diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java index d4a92f103..b406030f0 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorTests.java @@ -350,7 +350,7 @@ public void testProcessResponse_successful() throws Exception { Map knnMap = processor.buildMapWithKnnKeyAndOriginalValue(ingestDocument); List> modelTensorList = createMockVectorResult(); - processor.appendVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList); + processor.setVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList); assertEquals(12, ingestDocument.getSourceAndMetadata().size()); } @@ -398,6 +398,20 @@ public void testBuildVectorOutput_withNestedMap_successful() { assertNotNull(actionGamesKnn); } + public void test_updateDocument_appendVectorFieldsToDocument_successful() { + Map config = createPlainStringConfiguration(); + IngestDocument ingestDocument = createPlainIngestDocument(); + TextEmbeddingProcessor processor = createInstanceWithNestedMapConfiguration(config); + Map knnMap = processor.buildMapWithKnnKeyAndOriginalValue(ingestDocument); + List> modelTensorList = createMockVectorResult(); + processor.setVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList); + + List> modelTensorList1 = createMockVectorResult(); + processor.setVectorFieldsToDocument(ingestDocument, knnMap, modelTensorList1); + assertEquals(12, ingestDocument.getSourceAndMetadata().size()); + assertEquals(2, ((List) ingestDocument.getSourceAndMetadata().get("oriKey6_knn")).size()); + } + private List> createMockVectorResult() { List> modelTensorList = new ArrayList<>(); List number1 = ImmutableList.of(1.234f, 2.354f);