Add bwc test.

Signed-off-by: conggguan <[email protected]>
opensearch-project · Apr 29, 2024 · 6c6ac13 · 6c6ac13
1 parent 5ede643
commit 6c6ac13
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 2 deletions.
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
@@ -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<String, Object> twoPhaseResponse = search(getIndexNameForTest(), twoPhaseSparseEncodingQueryBuilder, 1);
+        Map<String, Object> firstInnerHit = getFirstInnerHit(twoPhaseResponse);
+        assertEquals("0", firstInnerHit.get("_id"));
+        float minExpectedScore = computeExpectedScore(modelId, testRankFeaturesDoc1, TEXT_1);
+        assertTrue(minExpectedScore <= objectToFloat(firstInnerHit.get("_score")));
+    }
+}
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
@@ -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<String, Object> twoPhaseResponse = search(getIndexNameForTest(), twoPhaseSparseEncodingQueryBuilder, 1);
+        Map<String, Object> 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
@@ -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;