[Backport 2.x] add support for builder constructor in neural query bu…

…ilder (#1084) * add support for builder constructor in neural query builder (#1047) Signed-off-by: will-hwang <[email protected]>
opensearch-project · Jan 9, 2025 · 452db42 · 452db42
1 parent 0b84f31
commit 452db42
Show file tree

Hide file tree

Showing 25 changed files with 604 additions and 548 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Support new knn query parameter expand_nested ([#1013](https://github.com/opensearch-project/neural-search/pull/1013))
 - Support empty string for fields in text embedding processor ([#1041](https://github.com/opensearch-project/neural-search/pull/1041))
 - Optimize ML inference connection retry logic ([#1054](https://github.com/opensearch-project/neural-search/pull/1054))
+- Support for builder constructor in Neural Query Builder ([#1047](https://github.com/opensearch-project/neural-search/pull/1047))
 ### Bug Fixes
 - Address inconsistent scoring in hybrid query results ([#998](https://github.com/opensearch-project/neural-search/pull/998))
 - Fix bug where ingested document has list of nested objects ([#1040](https://github.com/opensearch-project/neural-search/pull/1040))

diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java
@@ -122,11 +122,12 @@ private HybridQueryBuilder getQueryBuilder(
         final Map<String, ?> methodParameters,
         final RescoreContext rescoreContext
     ) {
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
-        neuralQueryBuilder.fieldName("passage_embedding");
-        neuralQueryBuilder.modelId(modelId);
-        neuralQueryBuilder.queryText(QUERY);
-        neuralQueryBuilder.k(5);
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .modelId(modelId)
+            .queryText(QUERY)
+            .k(5)
+            .build();
         if (expandNestedDocs != null) {
             neuralQueryBuilder.expandNested(expandNestedDocs);
         }

diff --git a/...tart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchWithRescoreIT.java b/...tart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchWithRescoreIT.java
@@ -93,11 +93,12 @@ private HybridQueryBuilder getQueryBuilder(
         final Map<String, ?> methodParameters,
         final RescoreContext rescoreContextForNeuralQuery
     ) {
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
-        neuralQueryBuilder.fieldName(VECTOR_EMBEDDING_FIELD);
-        neuralQueryBuilder.modelId(modelId);
-        neuralQueryBuilder.queryText(QUERY);
-        neuralQueryBuilder.k(5);
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName(VECTOR_EMBEDDING_FIELD)
+            .modelId(modelId)
+            .queryText(QUERY)
+            .k(5)
+            .build();
         if (methodParameters != null) {
             neuralQueryBuilder.methodParameters(methodParameters);
         }

diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java
@@ -51,37 +51,25 @@ public void testKnnRadialSearch_E2EFlow() throws Exception {
     }
 
     private void validateIndexQuery(final String modelId) {
-        NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = new NeuralQueryBuilder(
-            "passage_embedding",
-            TEXT,
-            TEST_IMAGE_TEXT,
-            modelId,
-            null,
-            null,
-            0.01f,
-            null,
-            null,
-            null,
-            null,
-            null
-        );
+        NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .queryText(TEXT)
+            .queryImage(TEST_IMAGE_TEXT)
+            .modelId(modelId)
+            .minScore(0.01f)
+            .build();
+
         Map<String, Object> responseWithMinScoreQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1);
         assertNotNull(responseWithMinScoreQuery);
 
-        NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder(
-            "passage_embedding",
-            TEXT,
-            TEST_IMAGE_TEXT,
-            modelId,
-            null,
-            100000f,
-            null,
-            null,
-            null,
-            null,
-            null,
-            null
-        );
+        NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .queryText(TEXT)
+            .queryImage(TEST_IMAGE_TEXT)
+            .modelId(modelId)
+            .maxDistance(100000f)
+            .build();
+
         Map<String, Object> responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1);
         assertNotNull(responseWithMaxDistanceQuery);
     }

diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java
@@ -53,20 +53,13 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception {
     private void validateTestIndex(final String modelId) throws Exception {
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(2, docCount);
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
-            "passage_embedding",
-            TEXT,
-            TEST_IMAGE_TEXT,
-            modelId,
-            1,
-            null,
-            null,
-            null,
-            null,
-            null,
-            null,
-            null
-        );
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .queryText(TEXT)
+            .queryImage(TEST_IMAGE_TEXT)
+            .modelId(modelId)
+            .k(1)
+            .build();
         Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
         assertNotNull(response);
     }

diff --git a/...upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java b/...upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java
@@ -76,8 +76,14 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws
 
     public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Exception {
         waitForClusterHealthGreen(NODES_BWC_CLUSTER);
-        NeuralQueryBuilder neuralQueryBuilderWithoutModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1);
-        NeuralQueryBuilder neuralQueryBuilderWithModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1);
+        NeuralQueryBuilder neuralQueryBuilderWithoutModelId = NeuralQueryBuilder.builder()
+            .fieldName(TEST_ENCODING_FIELD)
+            .queryText(TEXT_1)
+            .build();
+        NeuralQueryBuilder neuralQueryBuilderWithModelId = NeuralQueryBuilder.builder()
+            .fieldName(TEST_ENCODING_FIELD)
+            .queryText(TEXT_1)
+            .build();
 
         if (isRunningAgainstOldCluster()) {
             String modelId = uploadTextEmbeddingModel();

diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java
@@ -51,11 +51,12 @@ public void testTextEmbeddingProcessor_E2EFlow() throws Exception {
     private void validateTestIndex(final String modelId) throws Exception {
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(2, docCount);
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
-        neuralQueryBuilder.fieldName("passage_embedding");
-        neuralQueryBuilder.modelId(modelId);
-        neuralQueryBuilder.queryText(TEXT);
-        neuralQueryBuilder.k(1);
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .queryText(TEXT)
+            .modelId(modelId)
+            .k(1)
+            .build();
         Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
         assertNotNull(response);
     }

diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchIT.java
@@ -126,11 +126,12 @@ private HybridQueryBuilder getQueryBuilder(
         final Map<String, ?> methodParameters,
         final RescoreContext rescoreContextForNeuralQuery
     ) {
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
-        neuralQueryBuilder.fieldName(VECTOR_EMBEDDING_FIELD);
-        neuralQueryBuilder.modelId(modelId);
-        neuralQueryBuilder.queryText(QUERY);
-        neuralQueryBuilder.k(5);
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName(VECTOR_EMBEDDING_FIELD)
+            .modelId(modelId)
+            .queryText(QUERY)
+            .k(5)
+            .build();
         if (expandNestedDocs != null) {
             neuralQueryBuilder.expandNested(expandNestedDocs);
         }

diff --git a/...ling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchWithRescoreIT.java b/...ling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/HybridSearchWithRescoreIT.java
@@ -127,11 +127,12 @@ private HybridQueryBuilder getQueryBuilder(
         final Map<String, ?> methodParameters,
         final RescoreContext rescoreContextForNeuralQuery
     ) {
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
-        neuralQueryBuilder.fieldName(VECTOR_EMBEDDING_FIELD);
-        neuralQueryBuilder.modelId(modelId);
-        neuralQueryBuilder.queryText(QUERY);
-        neuralQueryBuilder.k(5);
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName(VECTOR_EMBEDDING_FIELD)
+            .modelId(modelId)
+            .queryText(QUERY)
+            .k(5)
+            .build();
         if (methodParameters != null) {
             neuralQueryBuilder.methodParameters(methodParameters);
         }

diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java
@@ -77,37 +77,25 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo
         assertEquals(numberOfDocs, docCount);
         loadModel(modelId);
 
-        NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = new NeuralQueryBuilder(
-            "passage_embedding",
-            text,
-            imageText,
-            modelId,
-            null,
-            null,
-            0.01f,
-            null,
-            null,
-            null,
-            null,
-            null
-        );
+        NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .queryText(text)
+            .queryImage(imageText)
+            .modelId(modelId)
+            .minScore(0.01f)
+            .build();
+
         Map<String, Object> responseWithMinScore = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1);
         assertNotNull(responseWithMinScore);
 
-        NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder(
-            "passage_embedding",
-            text,
-            imageText,
-            modelId,
-            null,
-            100000f,
-            null,
-            null,
-            null,
-            null,
-            null,
-            null
-        );
+        NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .queryText(text)
+            .queryImage(imageText)
+            .modelId(modelId)
+            .maxDistance(100000f)
+            .build();
+
         Map<String, Object> responseWithMaxScore = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1);
         assertNotNull(responseWithMaxScore);
     }

diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java
@@ -76,20 +76,14 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(numberOfDocs, docCount);
         loadModel(modelId);
-        NeuralQueryBuilder neuralQueryBuilderWithKQuery = new NeuralQueryBuilder(
-            "passage_embedding",
-            text,
-            imageText,
-            modelId,
-            1,
-            null,
-            null,
-            null,
-            null,
-            null,
-            null,
-            null
-        );
+        NeuralQueryBuilder neuralQueryBuilderWithKQuery = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .queryText(text)
+            .queryImage(imageText)
+            .modelId(modelId)
+            .k(1)
+            .build();
+
         Map<String, Object> responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1);
         assertNotNull(responseWithKQuery);
     }

diff --git a/...upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java b/...upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralQueryEnricherProcessorIT.java
@@ -94,8 +94,14 @@ public void testNeuralQueryEnricherProcessor_NeuralSparseSearch_E2EFlow() throws
     // the feature is introduced from 2.11
     public void testNeuralQueryEnricherProcessor_NeuralSearch_E2EFlow() throws Exception {
         waitForClusterHealthGreen(NODES_BWC_CLUSTER);
-        NeuralQueryBuilder neuralQueryBuilderWithoutModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1);
-        NeuralQueryBuilder neuralQueryBuilderWithModelId = new NeuralQueryBuilder().fieldName(TEST_ENCODING_FIELD).queryText(TEXT_1);
+        NeuralQueryBuilder neuralQueryBuilderWithoutModelId = NeuralQueryBuilder.builder()
+            .fieldName(TEST_ENCODING_FIELD)
+            .queryText(TEXT_1)
+            .build();
+        NeuralQueryBuilder neuralQueryBuilderWithModelId = NeuralQueryBuilder.builder()
+            .fieldName(TEST_ENCODING_FIELD)
+            .queryText(TEXT_1)
+            .build();
 
         switch (getClusterType()) {
             case OLD:

diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/SemanticSearchIT.java
@@ -71,11 +71,12 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(numberOfDocs, docCount);
         loadModel(modelId);
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder();
-        neuralQueryBuilder.fieldName("passage_embedding");
-        neuralQueryBuilder.modelId(modelId);
-        neuralQueryBuilder.queryText(text);
-        neuralQueryBuilder.k(1);
+        NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
+            .fieldName("passage_embedding")
+            .modelId(modelId)
+            .queryText(text)
+            .k(1)
+            .build();
         Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
         assertNotNull(response);
     }