Resolve comment

Signed-off-by: Junqiu Lei <[email protected]>
opensearch-project · Apr 22, 2024 · 390e6e8 · 390e6e8
1 parent 4f540fd
commit 390e6e8
Show file tree

Hide file tree

Showing 7 changed files with 80 additions and 20 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,7 +5,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 3.0](https://github.com/opensearch-project/neural-search/compare/2.x...HEAD)
 ### Features
-- Support k-NN radial search parameters in neural search([#697](https://github.com/opensearch-project/neural-search/pull/697))
 ### Enhancements
 ### Bug Fixes
 - Fix async actions are left in neural_sparse query ([#438](https://github.com/opensearch-project/neural-search/pull/438))
@@ -18,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 2.x](https://github.com/opensearch-project/neural-search/compare/2.13...2.x)
 ### Features
+- Support k-NN radial search parameters in neural search([#697](https://github.com/opensearch-project/neural-search/pull/697))
 ### Enhancements
 - BWC tests for text chunking processor ([#661](https://github.com/opensearch-project/neural-search/pull/661))
 - Allowing execution of hybrid query on index alias with filters ([#670](https://github.com/opensearch-project/neural-search/pull/670))

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -31,8 +31,8 @@ To send us a pull request, please:
 
 1. Fork the repository.
 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change.
-3. Include tests that check your new feature or bug fix. Ideally, we're looking for unit, integration, and BWC tests, but that depends on how big and critical your change is. 
-If you're adding an integration test and it is using local ML models, please make sure that the number of model deployments is limited, and you're using the smallest possible model. 
+3. Include tests that check your new feature or bug fix. Ideally, we're looking for unit, integration, and BWC tests, but that depends on how big and critical your change is.
+If you're adding an integration test and it is using local ML models, please make sure that the number of model deployments is limited, and you're using the smallest possible model.
 Each model deployment consumes resources, and having too many models may cause unexpected test failures.
 4. Ensure local tests pass.
 5. Commit to your fork using clear commit messages.

diff --git a/...ade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java b/...ade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java
@@ -4,11 +4,9 @@
  */
 package org.opensearch.neuralsearch.bwc;
 
-import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Locale;
-import java.util.Objects;
 import java.util.Optional;
 import org.junit.Before;
 import org.opensearch.common.settings.Settings;

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
@@ -50,10 +50,10 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception {
         }
     }
 
-    private void validateTestIndex(final String modelId) throws Exception {
+    private void validateTestIndex(final String modelId) {
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(2, docCount);
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
+        NeuralQueryBuilder neuralQueryBuilderWithKQuery = new NeuralQueryBuilder(
             "passage_embedding",
             TEXT,
             TEST_IMAGE_TEXT,
@@ -64,8 +64,35 @@ private void validateTestIndex(final String modelId) throws Exception {
             null,
             null
         );
-        Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
-        assertNotNull(response);
-    }
+        Map<String, Object> responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1);
+        assertNotNull(responseWithKQuery);
+
+        NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = new NeuralQueryBuilder(
+            "passage_embedding",
+            TEXT,
+            TEST_IMAGE_TEXT,
+            modelId,
+            null,
+            null,
+            0.01f,
+            null,
+            null
+        );
+        Map<String, Object> responseWithMinScoreQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1);
+        assertNotNull(responseWithMinScoreQuery);
 
+        NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder(
+            "passage_embedding",
+            TEXT,
+            TEST_IMAGE_TEXT,
+            modelId,
+            null,
+            10000f,
+            null,
+            null,
+            null
+        );
+        Map<String, Object> responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1);
+        assertNotNull(responseWithMaxDistanceQuery);
+    }
 }
diff --git a/...upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java b/...upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java
@@ -4,11 +4,9 @@
  */
 package org.opensearch.neuralsearch.bwc;
 
-import java.net.URL;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Locale;
-import java.util.Objects;
 import java.util.Optional;
 import org.junit.Before;
 import org.opensearch.common.settings.Settings;

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,7 +76,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(numberOfDocs, docCount);
         loadModel(modelId);
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
+        NeuralQueryBuilder neuralQueryBuilderWithKQuery = new NeuralQueryBuilder(
             "passage_embedding",
             text,
             imageText,
@@ -87,7 +87,35 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod
             null,
             null
         );
-        Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
-        assertNotNull(response);
+        Map<String, Object> responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1);
+        assertNotNull(responseWithKQuery);
+
+        NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = new NeuralQueryBuilder(
+            "passage_embedding",
+            text,
+            imageText,
+            modelId,
+            null,
+            null,
+            0.01f,
+            null,
+            null
+        );
+        Map<String, Object> responseWithMinScore = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1);
+        assertNotNull(responseWithMinScore);
+
+        NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder(
+            "passage_embedding",
+            text,
+            imageText,
+            modelId,
+            null,
+            10000f,
+            null,
+            null,
+            null
+        );
+        Map<String, Object> responseWithMaxScore = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1);
+        assertNotNull(responseWithMaxScore);
     }
 }
diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java
@@ -120,6 +120,11 @@ public NeuralQueryBuilder(StreamInput in) throws IOException {
         } else {
             this.modelId = in.readString();
         }
+        if (isClusterOnOrAfterMinReqVersionForRadialSearch()) {
+            this.k = in.readOptionalInt();
+        } else {
+            this.k = in.readVInt();
+        }
         this.k = in.readVInt();
         this.filter = in.readOptionalNamedWriteable(QueryBuilder.class);
         if (isClusterOnOrAfterMinReqVersionForRadialSearch()) {
@@ -138,7 +143,11 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         } else {
             out.writeString(this.modelId);
         }
-        out.writeVInt(this.k);
+        if (isClusterOnOrAfterMinReqVersionForRadialSearch()) {
+            out.writeOptionalInt(this.k);
+        } else {
+            out.writeVInt(this.k);
+        }
         out.writeOptionalNamedWriteable(this.filter);
         if (isClusterOnOrAfterMinReqVersionForRadialSearch()) {
             out.writeOptionalFloat(this.maxDistance);
@@ -219,8 +228,8 @@ public static NeuralQueryBuilder fromXContent(XContentParser parser) throws IOEx
             requireValue(neuralQueryBuilder.modelId(), "Model ID must be provided for neural query");
         }
 
-        long queryCount = validateKNNQueryType(neuralQueryBuilder);
-        if (queryCount == 0) {
+        boolean queryTypeIsProvided = validateKNNQueryType(neuralQueryBuilder);
+        if (queryTypeIsProvided == false) {
             neuralQueryBuilder.k(DEFAULT_K);
         }
 
@@ -359,7 +368,7 @@ private static boolean isClusterOnOrAfterMinReqVersionForRadialSearch() {
         return NeuralSearchClusterUtil.instance().getClusterMinVersion().onOrAfter(MINIMAL_SUPPORTED_VERSION_RADIAL_SEARCH);
     }
 
-    private static int validateKNNQueryType(NeuralQueryBuilder neuralQueryBuilder) {
+    private static boolean validateKNNQueryType(NeuralQueryBuilder neuralQueryBuilder) {
         int queryCount = 0;
         if (neuralQueryBuilder.k() != null) {
             queryCount++;
@@ -373,6 +382,6 @@ private static int validateKNNQueryType(NeuralQueryBuilder neuralQueryBuilder) {
         if (queryCount > 1) {
             throw new IllegalArgumentException("Only one of k, max_distance, or min_score can be provided");
         }
-        return queryCount;
+        return queryCount == 1;
     }
 }