Add basic integration test for segment replication

Signed-off-by: Suraj Singh <[email protected]>
opensearch-project · Jan 31, 2024 · eb5fe85 · eb5fe85
1 parent fe592f5
commit eb5fe85
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/src/test/java/org/opensearch/knn/index/OpenSearchIT.java b/src/test/java/org/opensearch/knn/index/OpenSearchIT.java
@@ -54,6 +54,30 @@ public static void setUpClass() throws IOException {
         testData = new TestUtils.TestData(testIndexVectors.getPath(), testQueries.getPath());
     }
 
+    /**
+     * This test performs basic validation steps such as index, search in order to verify that kNN works with Segment
+     * replication.
+     * @throws Exception
+     */
+    public void testSegmentReplication() throws Exception {
+        String field1 = "field1";
+        String field2 = "field2";
+        createKnnIndex(INDEX_NAME, getKNNSegrepIndexSettings(), createKnnIndexMapping(Arrays.asList(field1, field2), Arrays.asList(2, 2)));
+
+        Float[] vector = { 6.0f, 7.0f };
+        addKnnDoc(INDEX_NAME, "1", Arrays.asList(field1, field2), Arrays.asList(vector, vector));
+
+        addKnnDoc(INDEX_NAME, "2", field1, vector);
+        addKnnDoc(INDEX_NAME, "3", field1, vector);
+        addKnnDoc(INDEX_NAME, "4", field1, vector);
+
+        addKnnDoc(INDEX_NAME, "5", field2, vector);
+        addKnnDoc(INDEX_NAME, "6", field2, vector);
+
+        refreshAllIndices();
+        assertEquals(6, getDocCount(INDEX_NAME));
+    }
+
     public void testEndToEnd() throws Exception {
         String indexName = "test-index-1";
         KNNEngine knnEngine1 = KNNEngine.NMSLIB;

diff --git a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java
@@ -11,12 +11,15 @@
 import org.apache.hc.core5.http.io.entity.EntityUtils;
 import org.apache.hc.core5.net.URIBuilder;
 import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.common.Strings;
 import org.opensearch.common.xcontent.XContentHelper;
 import org.opensearch.core.xcontent.DeprecationHandler;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.query.MatchAllQueryBuilder;
+import org.opensearch.indices.replication.common.ReplicationType;
 import org.opensearch.knn.index.query.KNNQueryBuilder;
 import org.opensearch.knn.index.KNNSettings;
 import org.opensearch.knn.index.SpaceType;
@@ -93,6 +96,7 @@
 import static org.opensearch.knn.TestUtils.QUERY_VALUE;
 import static org.opensearch.knn.TestUtils.computeGroundTruthValues;
 
+import static org.opensearch.knn.index.KNNSettings.KNN_INDEX;
 import static org.opensearch.knn.index.SpaceType.L2;
 import static org.opensearch.knn.index.memory.NativeMemoryCacheManager.GRAPH_COUNT;
 import static org.opensearch.knn.index.util.KNNEngine.FAISS;
@@ -612,6 +616,18 @@ protected Settings getKNNDefaultIndexSettings() {
         return Settings.builder().put("number_of_shards", 1).put("number_of_replicas", 0).put("index.knn", true).build();
     }
 
+    /**
+     * Return default segment replication enabled index settings
+     */
+    protected Settings getKNNSegrepIndexSettings() {
+        return Settings.builder()
+            .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+            .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
+            .put(IndexMetadata.SETTING_REPLICATION_TYPE, ReplicationType.SEGMENT)
+            .put(KNN_INDEX, true)
+            .build();
+    }
+
     /**
      * Get Stats from KNN Plugin
      */