From 743ad9a1a31e61030785f8e5ff130f56ce752353 Mon Sep 17 00:00:00 2001 From: Heemin Kim Date: Mon, 22 Jan 2024 09:53:41 -0800 Subject: [PATCH] Few changes in integration test (#1395) Signed-off-by: Heemin Kim (cherry picked from commit 6abec195426344aa06a5905d721fc1fb005e3257) --- .../index/AdvancedFilteringUseCasesIT.java | 6 ++-- .../opensearch/knn/index/NestedSearchIT.java | 35 +++++++++---------- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/src/test/java/org/opensearch/knn/index/AdvancedFilteringUseCasesIT.java b/src/test/java/org/opensearch/knn/index/AdvancedFilteringUseCasesIT.java index 9ce994a78..b559b5760 100644 --- a/src/test/java/org/opensearch/knn/index/AdvancedFilteringUseCasesIT.java +++ b/src/test/java/org/opensearch/knn/index/AdvancedFilteringUseCasesIT.java @@ -448,14 +448,16 @@ public void testFiltering_whenNonNestedKNNAndNestedFilterAndNonNestedFieldWithNe private void validateFilterSearch(final String query, final String engine) throws IOException { String response = EntityUtils.toString(performSearch(INDEX_NAME, query).getEntity()); // Validate number of documents returned as the expected number of documents - Assert.assertEquals("For engine " + engine + " : ", DOCUMENT_IN_RESPONSE, parseHits(response)); + Assert.assertEquals("For engine " + engine + ", hits: ", DOCUMENT_IN_RESPONSE, parseHits(response)); + Assert.assertEquals("For engine " + engine + ", totalSearchHits: ", k, parseTotalSearchHits(response)); if (KNNEngine.getEngine(engine) == KNNEngine.FAISS) { // Update the filter threshold to 0 to ensure that we are hitting ANN Search use case for FAISS updateIndexSettings(INDEX_NAME, Settings.builder().put(KNNSettings.ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD, 0)); response = EntityUtils.toString(performSearch(INDEX_NAME, query).getEntity()); // Validate number of documents returned as the expected number of documents - Assert.assertEquals("For engine " + engine + " with ANN search :", DOCUMENT_IN_RESPONSE, parseHits(response)); + Assert.assertEquals("For engine " + engine + ", hits with ANN search :", DOCUMENT_IN_RESPONSE, parseHits(response)); + Assert.assertEquals("For engine " + engine + ", totalSearchHits with ANN search :", k, parseTotalSearchHits(response)); } } diff --git a/src/test/java/org/opensearch/knn/index/NestedSearchIT.java b/src/test/java/org/opensearch/knn/index/NestedSearchIT.java index 7a81e2f45..e4d828a01 100644 --- a/src/test/java/org/opensearch/knn/index/NestedSearchIT.java +++ b/src/test/java/org/opensearch/knn/index/NestedSearchIT.java @@ -63,17 +63,16 @@ public final void cleanUp() { public void testNestedSearchWithLucene_whenKIsTwo_thenReturnTwoResults() { createKnnIndex(2, KNNEngine.LUCENE.getName()); - String doc1 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED) - .addVectors(FIELD_NAME_VECTOR, new Float[] { 1f, 1f }, new Float[] { 1f, 1f }) - .build(); - addKnnDoc(INDEX_NAME, "1", doc1); - - String doc2 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED) - .addVectors(FIELD_NAME_VECTOR, new Float[] { 2f, 2f }, new Float[] { 2f, 2f }) - .build(); - addKnnDoc(INDEX_NAME, "2", doc2); + int totalDocCount = 15; + for (int i = 0; i < totalDocCount; i++) { + String doc = NestedKnnDocBuilder.create(FIELD_NAME_NESTED) + .addVectors(FIELD_NAME_VECTOR, new Float[] { (float) i, (float) i }, new Float[] { (float) i, (float) i }) + .build(); + addKnnDoc(INDEX_NAME, String.valueOf(i), doc); + } refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME); Float[] queryVector = { 1f, 1f }; Response response = queryNestedField(INDEX_NAME, 2, queryVector); @@ -86,17 +85,16 @@ public void testNestedSearchWithLucene_whenKIsTwo_thenReturnTwoResults() { public void testNestedSearchWithFaiss_whenKIsTwo_thenReturnTwoResults() { createKnnIndex(2, KNNEngine.FAISS.getName()); - String doc1 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED) - .addVectors(FIELD_NAME_VECTOR, new Float[] { 1f, 1f }, new Float[] { 1f, 1f }) - .build(); - addKnnDoc(INDEX_NAME, "1", doc1); - - String doc2 = NestedKnnDocBuilder.create(FIELD_NAME_NESTED) - .addVectors(FIELD_NAME_VECTOR, new Float[] { 2f, 2f }, new Float[] { 2f, 2f }) - .build(); - addKnnDoc(INDEX_NAME, "2", doc2); + int totalDocCount = 15; + for (int i = 0; i < totalDocCount; i++) { + String doc = NestedKnnDocBuilder.create(FIELD_NAME_NESTED) + .addVectors(FIELD_NAME_VECTOR, new Float[] { (float) i, (float) i }, new Float[] { (float) i, (float) i }) + .build(); + addKnnDoc(INDEX_NAME, String.valueOf(i), doc); + } refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME); Float[] queryVector = { 1f, 1f }; Response response = queryNestedField(INDEX_NAME, 2, queryVector); @@ -148,6 +146,7 @@ public void testNestedSearchWithFaiss_whenDoingExactSearch_thenReturnCorrectResu addKnnDoc(INDEX_NAME, String.valueOf(i), doc); } refreshIndex(INDEX_NAME); + forceMergeKnnIndex(INDEX_NAME); // Make it as an exact search by setting the threshold larger than size of filteredIds(6) updateIndexSettings(INDEX_NAME, Settings.builder().put(KNNSettings.ADVANCED_FILTERED_EXACT_SEARCH_THRESHOLD, 100));