diff --git a/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java b/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java index dac40a51..69239faf 100644 --- a/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java +++ b/src/test/java/org/opensearch/knn/index/engine/CommonTestUtils.java @@ -25,6 +25,7 @@ import org.opensearch.core.xcontent.*; import org.opensearch.knn.KNNResult; import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.KNNVectorSimilarityFunction; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.VectorDataType; @@ -97,7 +98,12 @@ public class CommonTestUtils { public static final String NON_EXISTENT_INTEGER_FIELD_NAME = "nonexistent_int_field"; public static Settings getDefaultIndexSettings() { - return Settings.builder().put("number_of_shards", 1).put("number_of_replicas", 0).put(KNN_INDEX, true).build(); + return Settings.builder() + .put("number_of_shards", 1) + .put("number_of_replicas", 0) + .put(KNN_INDEX, true) + .put(KNNSettings.KNN_DERIVED_SOURCE_ENABLED, true) + .build(); } public static String createIndexMapping(int dimension, SpaceType spaceType, VectorDataType vectorDataType) throws IOException { diff --git a/src/test/java/org/opensearch/knn/index/engine/InternalKNNEngineTests.java b/src/test/java/org/opensearch/knn/index/engine/InternalKNNEngineTests.java index 4f1f3d7e..0f7e0779 100644 --- a/src/test/java/org/opensearch/knn/index/engine/InternalKNNEngineTests.java +++ b/src/test/java/org/opensearch/knn/index/engine/InternalKNNEngineTests.java @@ -6,10 +6,10 @@ package org.opensearch.knn.index.engine; import com.carrotsearch.randomizedtesting.annotations.ThreadLeakFilters; +import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope; import org.apache.hc.core5.http.io.entity.EntityUtils; import org.apache.lucene.codecs.perfield.PerFieldKnnVectorsFormat; import org.apache.lucene.index.SegmentReader; -import org.junit.Ignore; import org.junit.Test; import org.opensearch.client.Request; import org.opensearch.client.Response; @@ -51,7 +51,8 @@ */ @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 1) -@ThreadLeakFilters(defaultFilters = true, filters = { ThreadLeakFiltersForTests.class }) +@ThreadLeakFilters(filters = { ThreadLeakFiltersForTests.class }) +@ThreadLeakScope(ThreadLeakScope.Scope.NONE) public class InternalKNNEngineTests extends OpenSearchIntegTestCase { /** ** Enable the http client *** */ @@ -93,7 +94,6 @@ private void createKnnIndexMappingWithJVectorEngine(int dimension, SpaceType spa throws Exception { String mapping = CommonTestUtils.createIndexMapping(dimension, spaceType, vectorDataType); Settings indexSettings = CommonTestUtils.getDefaultIndexSettings(); - // indexSettings = Settings.builder().put(indexSettings).put(INDEX_USE_COMPOUND_FILE.getKey(), false).build(); createKnnIndex(INDEX_NAME, indexSettings, mapping); } @@ -515,9 +515,8 @@ public void testMixedBatchSizesForQuantization() throws Exception { * @throws Exception exception */ @Test - @Ignore public void testQuantizationWithOverQueryParameter() throws Exception { - int dimension = 512; + int dimension = 128; final SpaceType spaceType = SpaceType.L2; final RestClient restClient = getRestClient(); createKnnIndexMappingWithJVectorEngine(dimension, spaceType, VectorDataType.FLOAT); @@ -526,9 +525,9 @@ public void testQuantizationWithOverQueryParameter() throws Exception { int batchSize = DEFAULT_MINIMUM_BATCH_SIZE_FOR_QUANTIZATION * 2; final float[][] vectors = TestUtils.generateRandomVectors(batchSize, dimension); - final int totalDocs = vectors.length; + final int expectedTotalDocs = vectors.length; - logger.info("Adding batch of vectors with size {} that is expected to trigger quantization", totalDocs); + logger.info("Adding batch of vectors with size {} that is expected to trigger quantization", expectedTotalDocs); CommonTestUtils.bulkAddKnnDocs(restClient, INDEX_NAME, FIELD_NAME, vectors, batchSize, false); CommonTestUtils.flushIndex(restClient, INDEX_NAME); @@ -537,7 +536,6 @@ public void testQuantizationWithOverQueryParameter() throws Exception { CommonTestUtils.forceMergeKnnIndex(restClient, INDEX_NAME); // Verify the total document count - int expectedTotalDocs = vectors.length; assertEquals(expectedTotalDocs, CommonTestUtils.getDocCount(restClient, INDEX_NAME)); // Perform search and verify recall @@ -570,9 +568,9 @@ public void testQuantizationWithOverQueryParameter() throws Exception { assertEquals(Math.min(k, expectedTotalDocs), results.size()); // calculate recall - logger.info("Calculating recall"); - float recall = ((float) results.stream().filter(r -> expectedDocIds.contains(r.getDocId())).count()) / ((float) k); - assertTrue("Expected recall to be lower than 0.7 but got " + recall, recall <= 0.7); + logger.info("Calculating recall with low overquery"); + float recallLowOverquery = ((float) results.stream().filter(r -> expectedDocIds.contains(r.getDocId())).count()) / ((float) k); + logger.info("Recall with low overquery: " + recallLowOverquery); // 2. Search with a high-overquery factor logger.info("Searching with high overquery factor"); @@ -594,9 +592,19 @@ public void testQuantizationWithOverQueryParameter() throws Exception { assertEquals(Math.min(k, expectedTotalDocs), results.size()); // calculate recall - logger.info("Calculating recall"); - recall = ((float) results.stream().filter(r -> expectedDocIds.contains(r.getDocId())).count()) / ((float) k); - assertTrue("Expected recall to be at least 0.9 but got " + recall, recall >= 0.9); + logger.info("Calculating recall with high overquery"); + float recallHighOverquery = ((float) results.stream().filter(r -> expectedDocIds.contains(r.getDocId())).count()) / ((float) k); + logger.info("Recall with high overquery: " + recallHighOverquery); + + // Verify that high overquery significantly improves recall compared to low overquery + assertTrue( + "Expected high overquery recall (" + + recallHighOverquery + + ") to be significantly better than low overquery recall (" + + recallLowOverquery + + ")", + recallHighOverquery > recallLowOverquery + 0.05 + ); } @Test