diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/IndexingIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/IndexingIT.java index 6b0bab9ee..bf4c8f7ec 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/IndexingIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/IndexingIT.java @@ -6,11 +6,14 @@ package org.opensearch.knn.bwc; import org.junit.Assert; +import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.VectorDataType; import org.opensearch.knn.index.engine.KNNEngine; +import java.util.List; import java.util.Map; import static org.opensearch.knn.TestUtils.KNN_ALGO_PARAM_EF_CONSTRUCTION_MIN_VALUE; @@ -52,6 +55,8 @@ public void testKNNIndexDefaultLegacyFieldMapping() throws Exception { createKnnIndex(testIndex, getKNNDefaultIndexSettings(), createKnnIndexMapping(TEST_FIELD, DIMENSIONS)); addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS); } else { + // update index setting to allow build graph always since we test graph count that are loaded into memory + updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0)); validateKNNIndexingOnUpgrade(NUM_DOCS); } } @@ -275,13 +280,14 @@ public void testNoParametersOnUpgrade() throws Exception { // KNN indexing tests when the cluster is upgraded to latest version public void validateKNNIndexingOnUpgrade(int numOfDocs) throws Exception { + updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0)); + forceMergeKnnIndex(testIndex); QUERY_COUNT = numOfDocs; validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, QUERY_COUNT, K); - cleanUpCache(); + clearCache(List.of(testIndex)); DOC_ID = numOfDocs; addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS); QUERY_COUNT = QUERY_COUNT + NUM_DOCS; - validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, QUERY_COUNT, K); forceMergeKnnIndex(testIndex); validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, QUERY_COUNT, K); deleteKNNIndex(testIndex); diff --git a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java index 1f9460397..db3555a8d 100644 --- a/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java +++ b/qa/restart-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java @@ -6,10 +6,10 @@ package org.opensearch.knn.bwc; import org.opensearch.common.settings.Settings; +import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.SpaceType; import java.util.Collections; -import java.util.Locale; import static org.opensearch.knn.TestUtils.KNN_ALGO_PARAM_M_MIN_VALUE; import static org.opensearch.knn.TestUtils.KNN_ALGO_PARAM_EF_CONSTRUCTION_MIN_VALUE; @@ -35,6 +35,8 @@ public void testKNNWarmupDefaultLegacyFieldMapping() throws Exception { createKnnIndex(testIndex, getKNNDefaultIndexSettings(), createKnnIndexMapping(TEST_FIELD, DIMENSIONS)); addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS); } else { + // update index setting to allow build graph always since we test graph count that are loaded into memory + updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0)); validateKNNWarmupOnUpgrade(); } } @@ -66,6 +68,8 @@ public void testKNNWarmupDefaultMethodFieldMapping() throws Exception { createKnnIndex(testIndex, getKNNDefaultIndexSettings(), createKNNIndexMethodFieldMapping(TEST_FIELD, DIMENSIONS)); addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS); } else { + // update index setting to allow build graph always since we test graph count that are loaded into memory + updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0)); validateKNNWarmupOnUpgrade(); } } @@ -86,22 +90,26 @@ public void testKNNWarmupCustomMethodFieldMapping() throws Exception { } public void validateKNNWarmupOnUpgrade() throws Exception { + // update index setting to allow build graph always since we test graph count that are loaded into memory + updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0)); int graphCount = getTotalGraphsInCache(); knnWarmup(Collections.singletonList(testIndex)); int totalGraph = getTotalGraphsInCache(); - assertTrue(String.format(Locale.ROOT, "[%d] is not greater than [%d]", totalGraph, graphCount), totalGraph > graphCount); + assertTrue(totalGraph > graphCount); QUERY_COUNT = NUM_DOCS; validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, QUERY_COUNT, K); DOC_ID = NUM_DOCS; addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, DOC_ID, NUM_DOCS); + forceMergeKnnIndex(testIndex); int updatedGraphCount = getTotalGraphsInCache(); knnWarmup(Collections.singletonList(testIndex)); assertTrue(getTotalGraphsInCache() > updatedGraphCount); QUERY_COUNT = QUERY_COUNT + NUM_DOCS; + validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, QUERY_COUNT, K); deleteKNNIndex(testIndex); } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ClearCacheIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ClearCacheIT.java index 24c674d0d..5ef3b1d97 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ClearCacheIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/ClearCacheIT.java @@ -25,6 +25,9 @@ public void testClearCache() throws Exception { createKnnIndex(testIndex, getKNNDefaultIndexSettings(), createKnnIndexMapping(TEST_FIELD, DIMENSIONS)); int docIdOld = 0; addKNNDocs(testIndex, TEST_FIELD, DIMENSIONS, docIdOld, NUM_DOCS); + int graphCount = getTotalGraphsInCache(); + knnWarmup(Collections.singletonList(testIndex)); + assertTrue(getTotalGraphsInCache() > graphCount); break; case UPGRADED: queryCnt = NUM_DOCS; @@ -42,14 +45,8 @@ public void testClearCache() throws Exception { // validation steps for Clear Cache API after upgrading all nodes from old version to new version public void validateClearCacheOnUpgrade(int queryCount) throws Exception { - int graphCount = getTotalGraphsInCache(); - knnWarmup(Collections.singletonList(testIndex)); - assertTrue(getTotalGraphsInCache() > graphCount); - validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, queryCount, K); - clearCache(Collections.singletonList(testIndex)); assertEquals(0, getTotalGraphsInCache()); - validateKNNSearch(testIndex, TEST_FIELD, DIMENSIONS, queryCount, K); } } diff --git a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java index 9cbb99d87..c760fde2a 100644 --- a/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java +++ b/qa/rolling-upgrade/src/test/java/org/opensearch/knn/bwc/WarmupIT.java @@ -5,6 +5,9 @@ package org.opensearch.knn.bwc; +import org.opensearch.common.settings.Settings; +import org.opensearch.knn.index.KNNSettings; + import java.util.Collections; import static org.opensearch.knn.TestUtils.NODES_BWC_CLUSTER; @@ -33,6 +36,7 @@ public void testKNNWarmup() throws Exception { docIdMixed = 2 * NUM_DOCS; totalDocsCountMixed = 3 * NUM_DOCS; } + updateIndexSettings(testIndex, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0)); validateKNNWarmupOnUpgrade(totalDocsCountMixed, docIdMixed); break; case UPGRADED: diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsFormat.java b/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsFormat.java index 520a9838d..dd326123e 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsFormat.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/NativeEngines990KnnVectorsFormat.java @@ -37,6 +37,10 @@ public NativeEngines990KnnVectorsFormat() { this(new Lucene99FlatVectorsFormat(new DefaultFlatVectorScorer())); } + public NativeEngines990KnnVectorsFormat(int approximateThreshold) { + this(new Lucene99FlatVectorsFormat(new DefaultFlatVectorScorer()), approximateThreshold); + } + public NativeEngines990KnnVectorsFormat(final FlatVectorsFormat flatVectorsFormat) { this(flatVectorsFormat, KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD_DEFAULT_VALUE); } diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/UnitTestCodec.java b/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/UnitTestCodec.java index 6998efbc3..d651f410a 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/UnitTestCodec.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN990Codec/UnitTestCodec.java @@ -21,6 +21,8 @@ * able to pick this class if its in test folder. Don't use this codec outside of testing */ public class UnitTestCodec extends FilterCodec { + private static final Integer BUILD_GRAPH_ALWAYS = 0; + public UnitTestCodec() { super("UnitTestCodec", KNNCodecVersion.current().getDefaultKnnCodecSupplier().get()); } @@ -30,7 +32,7 @@ public KnnVectorsFormat knnVectorsFormat() { return new PerFieldKnnVectorsFormat() { @Override public KnnVectorsFormat getKnnVectorsFormatForField(String field) { - return new NativeEngines990KnnVectorsFormat(); + return new NativeEngines990KnnVectorsFormat(BUILD_GRAPH_ALWAYS); } }; } diff --git a/src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java b/src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java index 06431bf07..587e80e5d 100644 --- a/src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java +++ b/src/test/java/org/opensearch/knn/KNNSingleNodeTestCase.java @@ -5,6 +5,7 @@ package org.opensearch.knn; +import org.opensearch.action.admin.indices.settings.put.UpdateSettingsRequest; import org.opensearch.core.action.ActionListener; import org.opensearch.core.common.bytes.BytesReference; import org.opensearch.cluster.ClusterName; @@ -14,6 +15,7 @@ import org.opensearch.cluster.block.ClusterBlocks; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.xcontent.XContentHelper; +import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.query.KNNQueryBuilder; import org.opensearch.knn.index.memory.NativeMemoryCacheManager; import org.opensearch.knn.index.memory.NativeMemoryLoadStrategy; @@ -104,6 +106,14 @@ protected void createKnnIndexMapping(String indexName, String fieldName, Integer OpenSearchAssertions.assertAcked(client().admin().indices().putMapping(request).actionGet()); } + /** + * Create simple k-NN mapping with engine + */ + protected void updateIndexSetting(String indexName, Settings setting) { + UpdateSettingsRequest request = new UpdateSettingsRequest(setting, indexName); + OpenSearchAssertions.assertAcked(client().admin().indices().updateSettings(request).actionGet()); + } + /** * Create simple k-NN mapping which can be nested. * e.g. fieldPath = "a.b.c" will create mapping for "c" as knn_vector @@ -140,6 +150,18 @@ protected Settings getKNNDefaultIndexSettings() { return Settings.builder().put("number_of_shards", 1).put("number_of_replicas", 0).put("index.knn", true).build(); } + /** + * Get default k-NN settings for test cases with build graph always + */ + protected Settings getKNNDefaultIndexSettingsBuildsGraphAlways() { + return Settings.builder() + .put("number_of_shards", 1) + .put("number_of_replicas", 0) + .put("index.knn", true) + .put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0) + .build(); + } + /** * Add a k-NN doc to an index */ diff --git a/src/test/java/org/opensearch/knn/index/KNNESSettingsTestIT.java b/src/test/java/org/opensearch/knn/index/KNNESSettingsTestIT.java index 36bf01f4e..20940f151 100644 --- a/src/test/java/org/opensearch/knn/index/KNNESSettingsTestIT.java +++ b/src/test/java/org/opensearch/knn/index/KNNESSettingsTestIT.java @@ -21,6 +21,9 @@ import static org.hamcrest.Matchers.containsString; public class KNNESSettingsTestIT extends KNNRestTestCase { + + public static final int ALWAYS_BUILD_GRAPH = 0; + /** * KNN Index writes should be blocked when the plugin disabled * @throws Exception Exception from test @@ -72,7 +75,7 @@ public void testQueriesPluginDisabled() throws Exception { } public void testItemRemovedFromCache_expiration() throws Exception { - createKnnIndex(INDEX_NAME, createKnnIndexMapping(FIELD_NAME, 2)); + createKnnIndex(INDEX_NAME, buildKNNIndexSettings(ALWAYS_BUILD_GRAPH), createKnnIndexMapping(FIELD_NAME, 2)); updateClusterSettings(KNNSettings.KNN_CACHE_ITEM_EXPIRY_ENABLED, true); updateClusterSettings(KNNSettings.KNN_CACHE_ITEM_EXPIRY_TIME_MINUTES, "1m"); diff --git a/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java b/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java index 18b9656e5..c25d2a390 100644 --- a/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java +++ b/src/test/java/org/opensearch/knn/index/KNNIndexShardTests.java @@ -14,6 +14,7 @@ import org.apache.lucene.util.StringHelper; import org.apache.lucene.util.Version; import org.mockito.Mockito; +import org.opensearch.common.settings.Settings; import org.opensearch.knn.KNNSingleNodeTestCase; import org.opensearch.index.IndexService; import org.opensearch.index.engine.Engine; @@ -71,6 +72,7 @@ public void testWarmup_emptyIndex() throws IOException { public void testWarmup_shardPresentInCache() throws InterruptedException, ExecutionException, IOException { IndexService indexService = createKNNIndex(testIndexName); createKnnIndexMapping(testIndexName, testFieldName, dimensions); + updateIndexSetting(testIndexName, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0).build()); addKnnDoc(testIndexName, "1", testFieldName, new Float[] { 2.5F, 3.5F }); searchKNNIndex(testIndexName, testFieldName, new float[] { 1.0f, 2.0f }, 1); @@ -85,6 +87,7 @@ public void testWarmup_shardPresentInCache() throws InterruptedException, Execut public void testWarmup_shardNotPresentInCache() throws InterruptedException, ExecutionException, IOException { IndexService indexService = createKNNIndex(testIndexName); createKnnIndexMapping(testIndexName, testFieldName, dimensions); + updateIndexSetting(testIndexName, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0).build()); IndexShard indexShard; KNNIndexShard knnIndexShard; @@ -106,6 +109,7 @@ public void testWarmup_shardNotPresentInCache() throws InterruptedException, Exe public void testGetAllEngineFileContexts() throws IOException, ExecutionException, InterruptedException { IndexService indexService = createKNNIndex(testIndexName); createKnnIndexMapping(testIndexName, testFieldName, dimensions); + updateIndexSetting(testIndexName, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0).build()); IndexShard indexShard = indexService.iterator().next(); KNNIndexShard knnIndexShard = new KNNIndexShard(indexShard); @@ -204,6 +208,7 @@ public void testClearCache_emptyIndex() { public void testClearCache_shardPresentInCache() { IndexService indexService = createKNNIndex(testIndexName); createKnnIndexMapping(testIndexName, testFieldName, dimensions); + updateIndexSetting(testIndexName, Settings.builder().put(KNNSettings.INDEX_KNN_ADVANCED_APPROXIMATE_THRESHOLD, 0).build()); addKnnDoc(testIndexName, String.valueOf(randomInt()), testFieldName, new Float[] { randomFloat(), randomFloat() }); IndexShard indexShard = indexService.iterator().next(); diff --git a/src/test/java/org/opensearch/knn/plugin/transport/ClearCacheTransportActionTests.java b/src/test/java/org/opensearch/knn/plugin/transport/ClearCacheTransportActionTests.java index 3222a3eb7..d79f2f738 100644 --- a/src/test/java/org/opensearch/knn/plugin/transport/ClearCacheTransportActionTests.java +++ b/src/test/java/org/opensearch/knn/plugin/transport/ClearCacheTransportActionTests.java @@ -33,7 +33,7 @@ public void testShardOperation() { KNNWarmupTransportAction knnWarmupTransportAction = node().injector().getInstance(KNNWarmupTransportAction.class); assertEquals(0, NativeMemoryCacheManager.getInstance().getIndicesCacheStats().size()); - IndexService indexService = createKNNIndex(testIndex); + IndexService indexService = createIndex(testIndex, getKNNDefaultIndexSettingsBuildsGraphAlways()); createKnnIndexMapping(testIndex, TEST_FIELD, DIMENSIONS); addKnnDoc(testIndex, String.valueOf(randomInt()), TEST_FIELD, new Float[] { randomFloat(), randomFloat() }); ShardRouting shardRouting = indexService.iterator().next().routingEntry(); diff --git a/src/test/java/org/opensearch/knn/plugin/transport/KNNWarmupTransportActionTests.java b/src/test/java/org/opensearch/knn/plugin/transport/KNNWarmupTransportActionTests.java index 1f72f78a7..f4b52246d 100644 --- a/src/test/java/org/opensearch/knn/plugin/transport/KNNWarmupTransportActionTests.java +++ b/src/test/java/org/opensearch/knn/plugin/transport/KNNWarmupTransportActionTests.java @@ -37,7 +37,7 @@ public void testShardOperation() throws IOException, ExecutionException, Interru KNNWarmupTransportAction knnWarmupTransportAction = node().injector().getInstance(KNNWarmupTransportAction.class); assertEquals(0, NativeMemoryCacheManager.getInstance().getIndicesCacheStats().size()); - indexService = createKNNIndex(testIndexName); + indexService = createIndex(testIndexName, getKNNDefaultIndexSettingsBuildsGraphAlways()); createKnnIndexMapping(testIndexName, testFieldName, dimensions); shardRouting = indexService.iterator().next().routingEntry(); diff --git a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java index 2c2efed3d..93f6b0792 100644 --- a/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java +++ b/src/testFixtures/java/org/opensearch/knn/KNNRestTestCase.java @@ -1309,6 +1309,7 @@ public void addKNNDocs(String testIndex, String testField, int dimension, int fi Arrays.fill(indexVector, (float) i); addKnnDoc(testIndex, Integer.toString(i), testField, indexVector); } + flushIndex(testIndex); } public void addKNNByteDocs(String testIndex, String testField, int dimension, int firstDocID, int numDocs) throws IOException { @@ -1317,6 +1318,7 @@ public void addKNNByteDocs(String testIndex, String testField, int dimension, in Arrays.fill(indexVector, (byte) i); addKnnDoc(testIndex, Integer.toString(i), testField, indexVector); } + flushIndex(testIndex); } public void validateKNNSearch(String testIndex, String testField, int dimension, int numDocs, int k) throws Exception { @@ -1792,6 +1794,13 @@ protected void refreshIndex(final String index) throws IOException { assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); } + protected void flushIndex(final String index) throws IOException { + Request request = new Request("POST", "/" + index + "/_flush"); + + Response response = client().performRequest(request); + assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + protected void addKnnDocWithAttributes(String docId, float[] vector, Map fieldValues) throws IOException { Request request = new Request("POST", "/" + INDEX_NAME + "/_doc/" + docId + "?refresh=true");