From 390e6e827228edd3c44b2e97f93285c5648f7639 Mon Sep 17 00:00:00 2001 From: Junqiu Lei Date: Mon, 22 Apr 2024 11:19:17 -0700 Subject: [PATCH] Resolve comment Signed-off-by: Junqiu Lei --- CHANGELOG.md | 2 +- CONTRIBUTING.md | 4 +- .../AbstractRestartUpgradeRestTestCase.java | 2 - .../neuralsearch/bwc/MultiModalSearchIT.java | 37 ++++++++++++++++--- .../bwc/AbstractRollingUpgradeTestCase.java | 2 - .../neuralsearch/bwc/MultiModalSearchIT.java | 34 +++++++++++++++-- .../query/NeuralQueryBuilder.java | 19 +++++++--- 7 files changed, 80 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 33b7e6613..6cb23f3a6 100644 --- 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 index 58562826d..67ba99e66 100644 --- 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/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java index 395573c6a..331ee714a 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRestartUpgradeRestTestCase.java +++ b/qa/restart-upgrade/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 index cbe210911..466460547 100644 --- 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 response = search(getIndexNameForTest(), neuralQueryBuilder, 1); - assertNotNull(response); - } + Map 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 responseWithMinScoreQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); + assertNotNull(responseWithMinScoreQuery); + NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder( + "passage_embedding", + TEXT, + TEST_IMAGE_TEXT, + modelId, + null, + 10000f, + null, + null, + null + ); + Map responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); + assertNotNull(responseWithMaxDistanceQuery); + } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java index ed1613e2f..489aefd6e 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/AbstractRollingUpgradeTestCase.java +++ b/qa/rolling-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 index e5b3f2b82..b757126e4 100644 --- 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 response = search(getIndexNameForTest(), neuralQueryBuilder, 1); - assertNotNull(response); + Map responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1); + assertNotNull(responseWithKQuery); + + NeuralQueryBuilder neuralQueryBuilderWithMinScoreQuery = new NeuralQueryBuilder( + "passage_embedding", + text, + imageText, + modelId, + null, + null, + 0.01f, + null, + null + ); + Map responseWithMinScore = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); + assertNotNull(responseWithMinScore); + + NeuralQueryBuilder neuralQueryBuilderWithMaxDistanceQuery = new NeuralQueryBuilder( + "passage_embedding", + text, + imageText, + modelId, + null, + 10000f, + null, + null, + null + ); + Map 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 index f016d128c..94d9b3481 100644 --- 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; } }