From eb5fe855e4b16d3a92981f78b7cdad3642ac21ab Mon Sep 17 00:00:00 2001 From: Suraj Singh Date: Mon, 5 Jun 2023 16:51:26 -0700 Subject: [PATCH] Add basic integration test for segment replication Signed-off-by: Suraj Singh --- .../opensearch/knn/index/OpenSearchIT.java | 24 +++++++++++++++++++ .../org/opensearch/knn/KNNRestTestCase.java | 16 +++++++++++++ 2 files changed, 40 insertions(+) diff --git a/src/test/java/org/opensearch/knn/index/OpenSearchIT.java b/src/test/java/org/opensearch/knn/index/OpenSearchIT.java index 4130e2df5..21d5f3550 100644 --- 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 index 1d984927f..663200b25 100644 --- 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 */