Radial search 2

CHANGELOG.md

@@ -19,6 +19,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))

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.

qa/restart-upgrade/build.gradle

@@ -90,6 +90,13 @@ task testAgainstOldCluster(type: StandaloneRestIntegTestTask) {
         }
     }
 
+    // Excluding the k-NN radial search tests because we introduce this feature in 2.14
+    if (ext.neural_search_bwc_version.startsWith("2.9") || ext.neural_search_bwc_version.startsWith("2.10") || ext.neural_search_bwc_version.startsWith("2.11") || ext.neural_search_bwc_version.startsWith("2.12") || ext.neural_search_bwc_version.startsWith("2.13")){
+        filter {
+            excludeTestsMatching "org.opensearch.neuralsearch.bwc.KnnRadialSearchIT.*"
+        }
+    }
+
     nonInputProperties.systemProperty('tests.rest.cluster', "${-> testClusters."${baseName}".allHttpSocketURI.join(",")}")
     nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
     systemProperty 'tests.security.manager', 'false'
@@ -139,6 +146,13 @@ task testAgainstNewCluster(type: StandaloneRestIntegTestTask) {
         }
     }
 
+    // Excluding the k-NN radial search tests because we introduce this feature in 2.14
+    if (ext.neural_search_bwc_version.startsWith("2.9") || ext.neural_search_bwc_version.startsWith("2.10") || ext.neural_search_bwc_version.startsWith("2.11") || ext.neural_search_bwc_version.startsWith("2.12") || ext.neural_search_bwc_version.startsWith("2.13")){
+        filter {
+            excludeTestsMatching "org.opensearch.neuralsearch.bwc.KnnRadialSearchIT.*"
+        }
+    }
+
     nonInputProperties.systemProperty('tests.rest.cluster', "${-> testClusters."${baseName}".allHttpSocketURI.join(",")}")
     nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
     systemProperty 'tests.security.manager', 'false'

qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.bwc;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Map;
+import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER;
+import static org.opensearch.neuralsearch.util.TestUtils.TEXT_IMAGE_EMBEDDING_PROCESSOR;
+import static org.opensearch.neuralsearch.util.TestUtils.getModelId;
+import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
+
+public class KnnRadialSearchIT extends AbstractRestartUpgradeRestTestCase {
+    private static final String PIPELINE_NAME = "radial-search-pipeline";
+    private static final String TEST_FIELD = "passage_text";
+    private static final String TEST_IMAGE_FIELD = "passage_image";
+    private static final String TEXT = "Hello world";
+    private static final String TEXT_1 = "Hello world a";
+    private static final String TEST_IMAGE_TEXT = "/9j/4AAQSkZJRgABAQAASABIAAD";
+    private static final String TEST_IMAGE_TEXT_1 = "/9j/4AAQSkZJRgbdwoeicfhoid";
+
+    // Test rolling-upgrade with kNN radial search
+    // Create Text Image Embedding Processor, Ingestion Pipeline and add document
+    // Validate radial query, pipeline and document count in restart-upgrade scenario
+    public void testKnnRadialSearch_E2EFlow() throws Exception {
+        waitForClusterHealthGreen(NODES_BWC_CLUSTER);
+
+        if (isRunningAgainstOldCluster()) {
+            String modelId = uploadTextEmbeddingModel();
+            loadModel(modelId);
+            createPipelineForTextImageProcessor(modelId, PIPELINE_NAME);
+            createIndexWithConfiguration(
+                getIndexNameForTest(),
+                Files.readString(Path.of(classLoader.getResource("processor/IndexMappingMultipleShard.json").toURI())),
+                PIPELINE_NAME
+            );
+            addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT);
+        } else {
+            String modelId = null;
+            try {
+                modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
+                loadModel(modelId);
+                addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_1, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_1);
+                validateIndexQuery(modelId);
+            } finally {
+                wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
+            }
+        }
+    }
+
+    private void validateIndexQuery(final String modelId) {
+        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,
+            100000f,
+            null,
+            null,
+            null
+        );
+        Map<String, Object> responseWithMaxDistanceQuery = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1);
+        assertNotNull(responseWithMaxDistanceQuery);
+    }
+}

qa/restart-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java

@@ -53,7 +53,17 @@ public void testTextImageEmbeddingProcessor_E2EFlow() throws Exception {
     private void validateTestIndex(final String modelId) throws Exception {
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(2, docCount);
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder("passage_embedding", TEXT, TEST_IMAGE_TEXT, modelId, 1, null, null);
+        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(
+            "passage_embedding",
+            TEXT,
+            TEST_IMAGE_TEXT,
+            modelId,
+            1,
+            null,
+            null,
+            null,
+            null
+        );
         Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
         assertNotNull(response);
     }

qa/rolling-upgrade/build.gradle

@@ -90,6 +90,13 @@ task testAgainstOldCluster(type: StandaloneRestIntegTestTask) {
         }
     }
 
+    // Excluding the k-NN radial search tests because we introduce this feature in 2.14
+    if (ext.neural_search_bwc_version.startsWith("2.9") || ext.neural_search_bwc_version.startsWith("2.10") || ext.neural_search_bwc_version.startsWith("2.11") || ext.neural_search_bwc_version.startsWith("2.12") || ext.neural_search_bwc_version.startsWith("2.13")){
+        filter {
+            excludeTestsMatching "org.opensearch.neuralsearch.bwc.KnnRadialSearchIT.*"
+        }
+    }
+
     nonInputProperties.systemProperty('tests.rest.cluster', "${-> testClusters."${baseName}".allHttpSocketURI.join(",")}")
     nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
     systemProperty 'tests.security.manager', 'false'
@@ -140,6 +147,13 @@ task testAgainstOneThirdUpgradedCluster(type: StandaloneRestIntegTestTask) {
         }
     }
 
+    // Excluding the k-NN radial search tests because we introduce this feature in 2.14
+    if (ext.neural_search_bwc_version.startsWith("2.9") || ext.neural_search_bwc_version.startsWith("2.10") || ext.neural_search_bwc_version.startsWith("2.11") || ext.neural_search_bwc_version.startsWith("2.12") || ext.neural_search_bwc_version.startsWith("2.13")){
+        filter {
+            excludeTestsMatching "org.opensearch.neuralsearch.bwc.KnnRadialSearchIT.*"
+        }
+    }
+
     nonInputProperties.systemProperty('tests.rest.cluster', "${-> testClusters."${baseName}".allHttpSocketURI.join(",")}")
     nonInputProperties.systemProperty('tests.clustername', "${-> testClusters."${baseName}".getName()}")
     systemProperty 'tests.security.manager', 'false'

qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/KnnRadialSearchIT.java

@@ -0,0 +1,108 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.bwc;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Map;
+import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER;
+import static org.opensearch.neuralsearch.util.TestUtils.TEXT_IMAGE_EMBEDDING_PROCESSOR;
+import static org.opensearch.neuralsearch.util.TestUtils.getModelId;
+import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
+
+public class KnnRadialSearchIT extends AbstractRollingUpgradeTestCase {
+    private static final String PIPELINE_NAME = "radial-search-pipeline";
+    private static final String TEST_FIELD = "passage_text";
+    private static final String TEST_IMAGE_FIELD = "passage_image";
+    private static final String TEXT = "Hello world";
+    private static final String TEXT_MIXED = "Hello world mixed";
+    private static final String TEXT_UPGRADED = "Hello world upgraded";
+    private static final String TEST_IMAGE_TEXT = "/9j/4AAQSkZJRgABAQAASABIAAD";
+    private static final String TEST_IMAGE_TEXT_MIXED = "/9j/4AAQSkZJRgbdwoeicfhoid";
+    private static final String TEST_IMAGE_TEXT_UPGRADED = "/9j/4AAQSkZJR8eydhgfwceocvlk";
+
+    private static final int NUM_DOCS_PER_ROUND = 1;
+    private static String modelId = "";
+
+    // Test rolling-upgrade with kNN radial search
+    // Create Text Image Embedding Processor, Ingestion Pipeline and add document
+    // Validate radial query, pipeline and document count in rolling-upgrade scenario
+    public void testKnnRadialSearch_E2EFlow() throws Exception {
+        waitForClusterHealthGreen(NODES_BWC_CLUSTER);
+        switch (getClusterType()) {
+            case OLD:
+                modelId = uploadTextImageEmbeddingModel();
+                loadModel(modelId);
+                createPipelineForTextImageProcessor(modelId, PIPELINE_NAME);
+                createIndexWithConfiguration(
+                    getIndexNameForTest(),
+                    Files.readString(Path.of(classLoader.getResource("processor/IndexMappings.json").toURI())),
+                    PIPELINE_NAME
+                );
+                addDocument(getIndexNameForTest(), "0", TEST_FIELD, TEXT, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT);
+                break;
+            case MIXED:
+                modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
+                int totalDocsCountMixed;
+                if (isFirstMixedRound()) {
+                    totalDocsCountMixed = NUM_DOCS_PER_ROUND;
+                    validateIndexQueryOnUpgrade(totalDocsCountMixed, modelId, TEXT, TEST_IMAGE_TEXT);
+                    addDocument(getIndexNameForTest(), "1", TEST_FIELD, TEXT_MIXED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_MIXED);
+                } else {
+                    totalDocsCountMixed = 2 * NUM_DOCS_PER_ROUND;
+                    validateIndexQueryOnUpgrade(totalDocsCountMixed, modelId, TEXT_MIXED, TEST_IMAGE_TEXT_MIXED);
+                }
+                break;
+            case UPGRADED:
+                try {
+                    modelId = getModelId(getIngestionPipeline(PIPELINE_NAME), TEXT_IMAGE_EMBEDDING_PROCESSOR);
+                    int totalDocsCountUpgraded = 3 * NUM_DOCS_PER_ROUND;
+                    loadModel(modelId);
+                    addDocument(getIndexNameForTest(), "2", TEST_FIELD, TEXT_UPGRADED, TEST_IMAGE_FIELD, TEST_IMAGE_TEXT_UPGRADED);
+                    validateIndexQueryOnUpgrade(totalDocsCountUpgraded, modelId, TEXT_UPGRADED, TEST_IMAGE_TEXT_UPGRADED);
+                } finally {
+                    wipeOfTestResources(getIndexNameForTest(), PIPELINE_NAME, modelId, null);
+                }
+                break;
+            default:
+                throw new IllegalStateException("Unexpected value: " + getClusterType());
+        }
+    }
+
+    private void validateIndexQueryOnUpgrade(final int numberOfDocs, final String modelId, final String text, final String imageText)
+        throws Exception {
+        int docCount = getDocCount(getIndexNameForTest());
+        assertEquals(numberOfDocs, docCount);
+        loadModel(modelId);
+
+        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,
+            100000f,
+            null,
+            null,
+            null
+        );
+        Map<String, Object> responseWithMaxScore = search(getIndexNameForTest(), neuralQueryBuilderWithMaxDistanceQuery, 1);
+        assertNotNull(responseWithMaxScore);
+    }
+}

qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/MultiModalSearchIT.java

@@ -76,8 +76,18 @@ private void validateTestIndexOnUpgrade(final int numberOfDocs, final String mod
         int docCount = getDocCount(getIndexNameForTest());
         assertEquals(numberOfDocs, docCount);
         loadModel(modelId);
-        NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder("passage_embedding", text, imageText, modelId, 1, null, null);
-        Map<String, Object> response = search(getIndexNameForTest(), neuralQueryBuilder, 1);
-        assertNotNull(response);
+        NeuralQueryBuilder neuralQueryBuilderWithKQuery = new NeuralQueryBuilder(
+            "passage_embedding",
+            text,
+            imageText,
+            modelId,
+            1,
+            null,
+            null,
+            null,
+            null
+        );
+        Map<String, Object> responseWithKQuery = search(getIndexNameForTest(), neuralQueryBuilderWithKQuery, 1);
+        assertNotNull(responseWithKQuery);
     }
 }