From 0cffa25dacdef1bccfc071aaf5714f91e0252176 Mon Sep 17 00:00:00 2001 From: Bo Zhang Date: Wed, 18 Dec 2024 13:44:19 -0800 Subject: [PATCH] [Backport] manually backport 1013 to 2.x (#1028) * Support of new k-NN query parameter expand_nested. Signed-off-by: Bo Zhang (cherry picked from commit fa149d43f98ae170b25d489c1639de6c00e0d606) * Remove mistakenly added code from HybridSearchIT. Signed-off-by: Bo Zhang (cherry picked from commit eaa7779932967feb49d4b9de2f86b51ff06c336a) --- CHANGELOG.md | 1 + .../neuralsearch/bwc/HybridSearchIT.java | 17 +++- .../neuralsearch/bwc/KnnRadialSearchIT.java | 2 + .../neuralsearch/bwc/MultiModalSearchIT.java | 1 + .../neuralsearch/bwc/HybridSearchIT.java | 46 ++++++---- .../neuralsearch/bwc/KnnRadialSearchIT.java | 2 + .../neuralsearch/bwc/MultiModalSearchIT.java | 1 + .../query/NeuralQueryBuilder.java | 15 ++++ .../processor/NormalizationProcessorIT.java | 3 + .../processor/ScoreCombinationIT.java | 60 ++++++++++++- .../processor/ScoreNormalizationIT.java | 90 +++++++++++++++++-- .../processor/TextEmbeddingProcessorIT.java | 1 + .../query/NeuralQueryBuilderTests.java | 9 ++ .../neuralsearch/query/NeuralQueryIT.java | 10 +++ 14 files changed, 230 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3be97248d..cc9960251 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Features ### Enhancements - Explainability in hybrid query ([#970](https://github.com/opensearch-project/neural-search/pull/970)) +- Support new knn query parameter expand_nested ([#1013](https://github.com/opensearch-project/neural-search/pull/1013)) ### Bug Fixes - Address inconsistent scoring in hybrid query results ([#998](https://github.com/opensearch-project/neural-search/pull/998)) ### Infrastructure 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 index fe69c577e..d4ae88a3f 100644 --- 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 @@ -12,6 +12,9 @@ import java.util.Map; import org.opensearch.index.query.MatchQueryBuilder; + +import static org.opensearch.knn.index.query.KNNQueryBuilder.EXPAND_NESTED_FIELD; +import static org.opensearch.neuralsearch.common.MinClusterVersionUtil.isClusterOnOrAfterMinReqVersion; import static org.opensearch.neuralsearch.util.TestUtils.getModelId; import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER; import static org.opensearch.neuralsearch.util.TestUtils.PARAM_NAME_WEIGHTS; @@ -71,9 +74,9 @@ private void validateNormalizationProcessor(final String fileName, final String modelId = getModelId(getIngestionPipeline(pipelineName), TEXT_EMBEDDING_PROCESSOR); loadModel(modelId); addDocuments(getIndexNameForTest(), false); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); - hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); + hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); validateTestIndex(getIndexNameForTest(), searchPipelineName, hybridQueryBuilder); } finally { wipeOfTestResources(getIndexNameForTest(), pipelineName, modelId, searchPipelineName); @@ -115,12 +118,20 @@ private void validateTestIndex(final String index, final String searchPipeline, } } - private HybridQueryBuilder getQueryBuilder(final String modelId, Map methodParameters, RescoreContext rescoreContext) { + private HybridQueryBuilder getQueryBuilder( + final String modelId, + final Boolean expandNestedDocs, + final Map methodParameters, + final RescoreContext rescoreContext + ) { NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); neuralQueryBuilder.fieldName("passage_embedding"); neuralQueryBuilder.modelId(modelId); neuralQueryBuilder.queryText(QUERY); neuralQueryBuilder.k(5); + if (expandNestedDocs != null) { + neuralQueryBuilder.expandNested(expandNestedDocs); + } 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 index d0994e711..5411caa0f 100644 --- 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 @@ -62,6 +62,7 @@ private void validateIndexQuery(final String modelId) { null, null, null, + null, null ); Map responseWithMinScoreQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); @@ -78,6 +79,7 @@ private void validateIndexQuery(final String modelId) { null, null, null, + null, null ); Map responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); 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 f35227041..bbca26cf6 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 @@ -64,6 +64,7 @@ private void validateTestIndex(final String modelId) throws Exception { null, null, null, + null, null ); Map response = search(getIndexNameForTest(), neuralQueryBuilder, 1); 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 index eeae7f7dd..78caca0e0 100644 --- 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 @@ -9,6 +9,8 @@ import java.util.Arrays; import java.util.List; import java.util.Map; +import java.util.Objects; + import org.opensearch.index.query.MatchQueryBuilder; import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER; import static org.opensearch.neuralsearch.util.TestUtils.PARAM_NAME_WEIGHTS; @@ -17,6 +19,8 @@ import static org.opensearch.neuralsearch.util.TestUtils.DEFAULT_COMBINATION_METHOD; import static org.opensearch.neuralsearch.util.TestUtils.getModelId; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilders; import org.opensearch.knn.index.query.rescore.RescoreContext; import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; @@ -31,6 +35,8 @@ public class HybridSearchIT extends AbstractRollingUpgradeTestCase { private static final String TEXT_UPGRADED = "Hi earth"; private static final String QUERY = "Hi world"; private static final int NUM_DOCS_PER_ROUND = 1; + private static final String VECTOR_EMBEDDING_FIELD = "passage_embedding"; + protected static final String RESCORE_QUERY = "hi"; private static String modelId = ""; // Test rolling-upgrade normalization processor when index with multiple shards @@ -61,13 +67,14 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr int totalDocsCountMixed; if (isFirstMixedRound()) { totalDocsCountMixed = NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, rescorer); addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, null, null); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); - validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + validateTestIndexOnUpgrade(totalDocsCountMixed, modelId, hybridQueryBuilder, null); } break; case UPGRADED: @@ -76,10 +83,11 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND; loadModel(modelId); addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, null, null); - HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder); - hybridQueryBuilder = getQueryBuilder(modelId, Map.of("ef_search", 100), RescoreContext.getDefault()); - validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder); + HybridQueryBuilder hybridQueryBuilder = getQueryBuilder(modelId, null, null, null); + QueryBuilder rescorer = QueryBuilders.matchQuery(TEST_FIELD, RESCORE_QUERY).boost(0.3f); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); + hybridQueryBuilder = getQueryBuilder(modelId, Boolean.FALSE, Map.of("ef_search", 100), RescoreContext.getDefault()); + validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId, hybridQueryBuilder, rescorer); } finally { wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, SEARCH_PIPELINE_NAME); } @@ -89,15 +97,19 @@ public void testNormalizationProcessor_whenIndexWithMultipleShards_E2EFlow() thr } } - private void validateTestIndexOnUpgrade(final int numberOfDocs, final String modelId, HybridQueryBuilder hybridQueryBuilder) - throws Exception { + private void validateTestIndexOnUpgrade( + final int numberOfDocs, + final String modelId, + HybridQueryBuilder hybridQueryBuilder, + QueryBuilder rescorer + ) throws Exception { int docCount = getDocCount(getIndexNameForTest()); assertEquals(numberOfDocs, docCount); loadModel(modelId); Map searchResponseAsMap = search( getIndexNameForTest(), hybridQueryBuilder, - null, + rescorer, 1, Map.of("search_pipeline", SEARCH_PIPELINE_NAME) ); @@ -112,19 +124,23 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod private HybridQueryBuilder getQueryBuilder( final String modelId, + final Boolean expandNestedDocs, final Map methodParameters, - final RescoreContext rescoreContext + final RescoreContext rescoreContextForNeuralQuery ) { NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); - neuralQueryBuilder.fieldName("passage_embedding"); + neuralQueryBuilder.fieldName(VECTOR_EMBEDDING_FIELD); neuralQueryBuilder.modelId(modelId); neuralQueryBuilder.queryText(QUERY); neuralQueryBuilder.k(5); + if (expandNestedDocs != null) { + neuralQueryBuilder.expandNested(expandNestedDocs); + } if (methodParameters != null) { neuralQueryBuilder.methodParameters(methodParameters); } - if (rescoreContext != null) { - neuralQueryBuilder.rescoreContext(rescoreContext); + if (Objects.nonNull(rescoreContextForNeuralQuery)) { + neuralQueryBuilder.rescoreContext(rescoreContextForNeuralQuery); } MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder("text", QUERY); 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 index 391e2135d..9dabc8d37 100644 --- 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 @@ -88,6 +88,7 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo null, null, null, + null, null ); Map responseWithMinScore = search(getIndexNameForTest(), neuralQueryBuilderWithMinScoreQuery, 1); @@ -104,6 +105,7 @@ private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String mo null, null, null, + null, null ); Map responseWithMaxScore = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1); 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 72976770d..2b1492319 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 @@ -87,6 +87,7 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod null, null, null, + null, null ); Map responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1); diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java index 915a79117..ce22507cb 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralQueryBuilder.java @@ -4,6 +4,7 @@ */ package org.opensearch.neuralsearch.query; +import static org.opensearch.knn.index.query.KNNQueryBuilder.EXPAND_NESTED_FIELD; import static org.opensearch.knn.index.query.KNNQueryBuilder.FILTER_FIELD; import static org.opensearch.knn.index.query.KNNQueryBuilder.MAX_DISTANCE_FIELD; import static org.opensearch.knn.index.query.KNNQueryBuilder.METHOD_PARAMS_FIELD; @@ -98,6 +99,7 @@ public static void initialize(MLCommonsClientAccessor mlClient) { private Integer k = null; private Float maxDistance = null; private Float minScore = null; + private Boolean expandNested; @VisibleForTesting @Getter(AccessLevel.PACKAGE) @Setter(AccessLevel.PACKAGE) @@ -132,6 +134,9 @@ public NeuralQueryBuilder(StreamInput in) throws IOException { this.maxDistance = in.readOptionalFloat(); this.minScore = in.readOptionalFloat(); } + if (isClusterOnOrAfterMinReqVersion(EXPAND_NESTED_FIELD.getPreferredName())) { + this.expandNested = in.readOptionalBoolean(); + } if (isClusterOnOrAfterMinReqVersion(METHOD_PARAMS_FIELD.getPreferredName())) { this.methodParameters = MethodParametersParser.streamInput(in, MinClusterVersionUtil::isClusterOnOrAfterMinReqVersion); } @@ -158,6 +163,9 @@ protected void doWriteTo(StreamOutput out) throws IOException { out.writeOptionalFloat(this.maxDistance); out.writeOptionalFloat(this.minScore); } + if (isClusterOnOrAfterMinReqVersion(EXPAND_NESTED_FIELD.getPreferredName())) { + out.writeOptionalBoolean(this.expandNested); + } if (isClusterOnOrAfterMinReqVersion(METHOD_PARAMS_FIELD.getPreferredName())) { MethodParametersParser.streamOutput(out, methodParameters, MinClusterVersionUtil::isClusterOnOrAfterMinReqVersion); } @@ -184,6 +192,9 @@ protected void doXContent(XContentBuilder xContentBuilder, Params params) throws if (Objects.nonNull(minScore)) { xContentBuilder.field(MIN_SCORE_FIELD.getPreferredName(), minScore); } + if (Objects.nonNull(expandNested)) { + xContentBuilder.field(EXPAND_NESTED_FIELD.getPreferredName(), expandNested); + } if (Objects.nonNull(methodParameters)) { MethodParametersParser.doXContent(xContentBuilder, methodParameters); } @@ -274,6 +285,8 @@ private static void parseQueryParams(XContentParser parser, NeuralQueryBuilder n neuralQueryBuilder.maxDistance(parser.floatValue()); } else if (MIN_SCORE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { neuralQueryBuilder.minScore(parser.floatValue()); + } else if (EXPAND_NESTED_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { + neuralQueryBuilder.expandNested(parser.booleanValue()); } else { throw new ParsingException( parser.getTokenLocation(), @@ -318,6 +331,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) { .filter(filter()) .maxDistance(maxDistance) .minScore(minScore) + .expandNested(expandNested) .k(k) .methodParameters(methodParameters) .rescoreContext(rescoreContext) @@ -346,6 +360,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) { k(), maxDistance(), minScore(), + expandNested(), vectorSetOnce::get, filter(), methodParameters(), diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java index 0fe9a77d6..385ea9428 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NormalizationProcessorIT.java @@ -98,6 +98,7 @@ public void testResultProcessor_whenOneShardAndQueryMatches_thenSuccessful() { null, null, null, + null, null ); TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); @@ -150,6 +151,7 @@ public void testResultProcessor_whenDefaultProcessorConfigAndQueryMatches_thenSu null, null, null, + null, null ); TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); @@ -191,6 +193,7 @@ public void testQueryMatches_whenMultipleShards_thenSuccessful() { null, null, null, + null, null ); TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java index 4ddf8d2c3..934fa94e0 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreCombinationIT.java @@ -224,7 +224,20 @@ public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessf HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); hybridQueryBuilderDefaultNorm.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderDefaultNorm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -249,7 +262,20 @@ public void testHarmonicMeanCombination_whenOneShardAndQueryMatches_thenSuccessf HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); hybridQueryBuilderL2Norm.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderL2Norm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -299,7 +325,20 @@ public void testGeometricMeanCombination_whenOneShardAndQueryMatches_thenSuccess HybridQueryBuilder hybridQueryBuilderDefaultNorm = new HybridQueryBuilder(); hybridQueryBuilderDefaultNorm.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderDefaultNorm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -324,7 +363,20 @@ public void testGeometricMeanCombination_whenOneShardAndQueryMatches_thenSuccess HybridQueryBuilder hybridQueryBuilderL2Norm = new HybridQueryBuilder(); hybridQueryBuilderL2Norm.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderL2Norm.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java index d851131f1..4ec0b49fa 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/ScoreNormalizationIT.java @@ -85,7 +85,20 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); hybridQueryBuilderArithmeticMean.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderArithmeticMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -110,7 +123,20 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); hybridQueryBuilderHarmonicMean.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderHarmonicMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -135,7 +161,20 @@ public void testL2Norm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); hybridQueryBuilderGeometricMean.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderGeometricMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -185,7 +224,20 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderArithmeticMean = new HybridQueryBuilder(); hybridQueryBuilderArithmeticMean.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderArithmeticMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -210,7 +262,20 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderHarmonicMean = new HybridQueryBuilder(); hybridQueryBuilderHarmonicMean.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderHarmonicMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); @@ -235,7 +300,20 @@ public void testMinMaxNorm_whenOneShardAndQueryMatches_thenSuccessful() { HybridQueryBuilder hybridQueryBuilderGeometricMean = new HybridQueryBuilder(); hybridQueryBuilderGeometricMean.add( - new NeuralQueryBuilder(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DOC_TEXT1, "", modelId, 5, null, null, null, null, null, null) + new NeuralQueryBuilder( + TEST_KNN_VECTOR_FIELD_NAME_1, + TEST_DOC_TEXT1, + "", + modelId, + 5, + null, + null, + null, + null, + null, + null, + null + ) ); hybridQueryBuilderGeometricMean.add(QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3)); diff --git a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java index 0ef9dda40..263a55273 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/TextEmbeddingProcessorIT.java @@ -125,6 +125,7 @@ public void testNestedFieldMapping_whenDocumentsIngested_thenSuccessful() throws null, null, null, + null, null ); QueryBuilder queryNestedLowerLevel = QueryBuilders.nestedQuery( diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java index 6d8e810f3..1d4a398cb 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryBuilderTests.java @@ -11,6 +11,7 @@ import static org.opensearch.core.xcontent.ToXContent.EMPTY_PARAMS; import static org.opensearch.index.query.AbstractQueryBuilder.BOOST_FIELD; import static org.opensearch.index.query.AbstractQueryBuilder.NAME_FIELD; +import static org.opensearch.knn.index.query.KNNQueryBuilder.EXPAND_NESTED_FIELD; import static org.opensearch.knn.index.query.KNNQueryBuilder.FILTER_FIELD; import static org.opensearch.knn.index.query.KNNQueryBuilder.MAX_DISTANCE_FIELD; import static org.opensearch.knn.index.query.KNNQueryBuilder.MIN_SCORE_FIELD; @@ -202,6 +203,7 @@ public void testFromXContent_whenBuiltWithOptionals_thenBuildSuccessfully() { "k": int, "boost": 10.0, "_name": "something", + "expandNestedDocs": true } } */ @@ -215,6 +217,7 @@ public void testFromXContent_whenBuiltWithOptionals_thenBuildSuccessfully() { .field(K_FIELD.getPreferredName(), K) .field(BOOST_FIELD.getPreferredName(), BOOST) .field(NAME_FIELD.getPreferredName(), QUERY_NAME) + .field(EXPAND_NESTED_FIELD.getPreferredName(), Boolean.TRUE) .endObject() .endObject(); @@ -229,6 +232,7 @@ public void testFromXContent_whenBuiltWithOptionals_thenBuildSuccessfully() { assertEquals(K, neuralQueryBuilder.k()); assertEquals(BOOST, neuralQueryBuilder.boost(), 0.0); assertEquals(QUERY_NAME, neuralQueryBuilder.queryName()); + assertEquals(Boolean.TRUE, neuralQueryBuilder.expandNested()); } @SneakyThrows @@ -428,6 +432,7 @@ public void testToXContent() { .modelId(MODEL_ID) .queryText(QUERY_TEXT) .k(K) + .expandNested(Boolean.TRUE) .filter(TEST_FILTER); XContentBuilder builder = XContentFactory.jsonBuilder(); @@ -454,6 +459,7 @@ public void testToXContent() { assertEquals(MODEL_ID, secondInnerMap.get(MODEL_ID_FIELD.getPreferredName())); assertEquals(QUERY_TEXT, secondInnerMap.get(QUERY_TEXT_FIELD.getPreferredName())); assertEquals(K, secondInnerMap.get(K_FIELD.getPreferredName())); + assertEquals(Boolean.TRUE, secondInnerMap.get(EXPAND_NESTED_FIELD.getPreferredName())); XContentBuilder xContentBuilder = XContentFactory.jsonBuilder(); assertEquals( xContentBuilderToMap(TEST_FILTER.toXContent(xContentBuilder, EMPTY_PARAMS)), @@ -718,6 +724,7 @@ public void testRewrite_whenVectorNull_thenReturnCopy() { .queryImage(IMAGE_TEXT) .modelId(MODEL_ID) .k(K) + .expandNested(Boolean.TRUE) .vectorSupplier(nullSupplier); QueryBuilder queryBuilder = neuralQueryBuilder.doRewrite(null); assertEquals(neuralQueryBuilder, queryBuilder); @@ -729,6 +736,7 @@ public void testRewrite_whenVectorSupplierAndVectorSet_thenReturnKNNQueryBuilder .queryImage(IMAGE_TEXT) .modelId(MODEL_ID) .k(K) + .expandNested(Boolean.TRUE) .methodParameters(Map.of("ef_search", 100)) .rescoreContext(RescoreContext.getDefault()) .vectorSupplier(TEST_VECTOR_SUPPLIER); @@ -739,6 +747,7 @@ public void testRewrite_whenVectorSupplierAndVectorSet_thenReturnKNNQueryBuilder .methodParameters(neuralQueryBuilder.methodParameters()) .rescoreContext(neuralQueryBuilder.rescoreContext()) .vector(TEST_VECTOR_SUPPLIER.get()) + .expandNested(Boolean.TRUE) .build(); QueryBuilder queryBuilder = neuralQueryBuilder.doRewrite(null); diff --git a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java index 210abd7ca..a7d349f07 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/NeuralQueryIT.java @@ -113,6 +113,7 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { null, null, null, + null, null ); @@ -135,6 +136,7 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { null, null, null, + null, Map.of("ef_search", 10), RescoreContext.getDefault() ); @@ -164,6 +166,7 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { null, null, null, + null, null ); @@ -194,6 +197,7 @@ public void testQueryWithBoostAndImageQueryAndRadialQuery() { null, null, null, + null, null ); @@ -250,6 +254,7 @@ public void testRescoreQuery() { null, null, null, + null, null ); @@ -329,6 +334,7 @@ public void testBooleanQuery_withMultipleNeuralQueries() { null, null, null, + null, null ); NeuralQueryBuilder neuralQueryBuilder2 = new NeuralQueryBuilder( @@ -342,6 +348,7 @@ public void testBooleanQuery_withMultipleNeuralQueries() { null, null, null, + null, null ); @@ -371,6 +378,7 @@ public void testBooleanQuery_withMultipleNeuralQueries() { null, null, null, + null, null ); @@ -428,6 +436,7 @@ public void testNestedQuery() { null, null, null, + null, null ); @@ -478,6 +487,7 @@ public void testFilterQuery() { null, null, null, + null, new MatchQueryBuilder("_id", "3"), null, null