diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java index 8ec54711a..b2ef70188 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java @@ -10,6 +10,7 @@ import java.util.Map; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.neuralsearch.query.NeuralSparseTwoPhaseParameters; import org.opensearch.neuralsearch.util.TestUtils; import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER; import static org.opensearch.neuralsearch.util.TestUtils.SPARSE_ENCODING_PROCESSOR; @@ -50,6 +51,7 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEST_TEXT_FIELD), List.of(TEXT_1) ); + } else { String modelId = null; try { @@ -64,6 +66,7 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEXT_2) ); validateTestIndex(modelId); + validateTwoPhaseQuery(modelId); } finally { wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); } @@ -87,4 +90,22 @@ private void validateTestIndex(final String modelId) throws Exception { float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc1, TEXT_1); assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score"))); } -} + + private void validateTwoPhaseQuery(final String modelId) throws Exception { + NeuralSparseQueryBuilder twoPhaseSparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_SPARSE_ENCODING_FIELD) + .queryText(TEXT_1) + .modelId(modelId) + .neuralSparseTwoPhaseParameters( + new NeuralSparseTwoPhaseParameters() + .enabled(true) + .window_size_expansion(2.0f) + .pruning_ratio(0.1f) + ); + + Map twoPhaseResponse = search(getIndexNameForTest(), twoPhaseSparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(twoPhaseResponse); + assertEquals("0", firstInnerHit.get("_id")); + float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc1, TEXT_1); + assertTrue(minExpectedScore <= objectToFloat(firstInnerHit.get("_score"))); + } +} \ No newline at end of file diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java index 0801ea201..f7b2e9c1c 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/NeuralSparseSearchIT.java @@ -10,6 +10,7 @@ import java.util.Map; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.neuralsearch.query.NeuralSparseTwoPhaseParameters; import org.opensearch.neuralsearch.util.TestUtils; import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER; import static org.opensearch.neuralsearch.util.TestUtils.SPARSE_ENCODING_PROCESSOR; @@ -57,6 +58,7 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEST_TEXT_FIELD), List.of(TEXT) ); + validateTwoPhaseQuery(modelId); break; case MIXED: modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), SPARSE_ENCODING_PROCESSOR); @@ -72,6 +74,7 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEST_TEXT_FIELD), List.of(TEXT_MIXED) ); + validateTwoPhaseQuery(modelId); } else { totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND; validateTestIndexOnUpgrade(totalDocsCountMixed, modelId); @@ -91,6 +94,7 @@ public void testSparseEncodingProcessor_E2EFlow() throws Exception { List.of(TEXT_UPGRADED) ); validateTestIndexOnUpgrade(totalDocsCountUpgraded, modelId); + validateTwoPhaseQuery(modelId); } finally { wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null); } @@ -117,4 +121,22 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc1, TEXT); assertTrue(minExpectedScore < objectToFloat(firstInnerHit.get("_score"))); } + + private void validateTwoPhaseQuery(final String modelId) throws Exception { + NeuralSparseQueryBuilder twoPhaseSparseEncodingQueryBuilder = new NeuralSparseQueryBuilder().fieldName(TEST_SPARSE_ENCODING_FIELD) + .queryText(TEXT+"c") + .modelId(modelId) + .neuralSparseTwoPhaseParameters( + new NeuralSparseTwoPhaseParameters() + .enabled(true) + .window_size_expansion(2.0f) + .pruning_ratio(0.1f) + ); + + Map twoPhaseResponse = search(getIndexNameForTest(), twoPhaseSparseEncodingQueryBuilder, 1); + Map firstInnerHit = getFirstInnerHit(twoPhaseResponse); + assertEquals("0", firstInnerHit.get("_id")); + float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc1, TEXT); + assertTrue(minExpectedScore <= objectToFloat(firstInnerHit.get("_score"))); + } } diff --git a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java index 47db9e412..3bd68a635 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java +++ b/src/main/java/org/opensearch/neuralsearch/query/NeuralSparseTwoPhaseParameters.java @@ -46,7 +46,7 @@ public class NeuralSparseTwoPhaseParameters implements Writeable { static final ParseField PRUNING_RATIO = new ParseField("pruning_ratio"); @VisibleForTesting static final ParseField ENABLED = new ParseField("enabled"); - private static final Version MINIMAL_SUPPORTED_VERSION_TWO_PHASE_SEARCH = Version.CURRENT; + private static final Version MINIMAL_SUPPORTED_VERSION_TWO_PHASE_SEARCH = Version.V_2_14_0; private Float window_size_expansion; private Float pruning_ratio; private Boolean enabled;