From 0553620d21dabd13814a87436009edf96789e4d9 Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Mon, 11 Sep 2023 11:50:06 -0700 Subject: [PATCH 01/10] Add graph stats to KNN Stats API Signed-off-by: Ryan Bogan --- .../KNN80Codec/KNN80DocValuesConsumer.java | 44 ++++++++++++++++++- .../knn/plugin/stats/KNNCounter.java | 8 +++- .../opensearch/knn/plugin/stats/KNNStats.java | 24 ++++++++++ .../knn/plugin/stats/StatNames.java | 4 +- 4 files changed, 76 insertions(+), 4 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index d9a30f75b..023607de4 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -9,6 +9,7 @@ import lombok.NonNull; import lombok.extern.log4j.Log4j2; import org.apache.lucene.store.ChecksumIndexInput; +import org.opensearch.common.StopWatch; import org.opensearch.common.xcontent.XContentHelper; import org.opensearch.core.common.bytes.BytesArray; import org.opensearch.core.xcontent.MediaTypeRegistry; @@ -74,9 +75,13 @@ class KNN80DocValuesConsumer extends DocValuesConsumer implements Closeable { @Override public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - delegatee.addBinaryField(field, valuesProducer); if (isKNNBinaryFieldRequired(field)) { + StopWatch stopWatch = new StopWatch(); + stopWatch.start(); addKNNBinaryField(field, valuesProducer); + stopWatch.stop(); + long time_in_millis = stopWatch.totalTime().millis(); + KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.set(KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.getCount() + time_in_millis); } } @@ -98,13 +103,19 @@ private KNNEngine getKNNEngine(@NonNull FieldInfo field) { } public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { - // Get values to be indexed + // Get values to be index BinaryDocValues values = valuesProducer.getBinary(field); KNNCodecUtil.Pair pair = KNNCodecUtil.getFloats(values); if (pair.vectors.length == 0 || pair.docs.length == 0) { logger.info("Skipping engine index creation as there are no vectors or docs in the documents"); return; } + KNNCounter.REFRESH_CURRENT_OPERATIONS.increment(); + KNNCounter.REFRESH_CURRENT_DOCS.set(KNNCounter.REFRESH_CURRENT_DOCS.getCount() + pair.docs.length); + KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.set( + KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getCount() + calculateArraySize(pair.vectors) + ); + // Increment counter for number of graph index requests KNNCounter.GRAPH_INDEX_REQUESTS.increment(); // Create library index either from model or from scratch @@ -135,6 +146,15 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) indexCreator = () -> createKNNIndexFromScratch(field, pair, knnEngine, indexPath); } + KNNCounter.REFRESH_CURRENT_OPERATIONS.set(KNNCounter.REFRESH_CURRENT_OPERATIONS.getCount() - 1); + KNNCounter.REFRESH_CURRENT_DOCS.set(KNNCounter.REFRESH_CURRENT_DOCS.getCount() - pair.docs.length); + KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.set( + KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getCount() - calculateArraySize(pair.vectors) + ); + KNNCounter.REFRESH_TOTAL_OPERATIONS.increment(); + System.out.println(KNNCounter.REFRESH_TOTAL_OPERATIONS.getCount()); + KNNCounter.REFRESH_TOTAL_DOCS.set(KNNCounter.REFRESH_TOTAL_DOCS.getCount() + pair.docs.length); + // This is a bit of a hack. We have to create an output here and then immediately close it to ensure that // engineFileName is added to the tracked files by Lucene's TrackingDirectoryWrapper. Otherwise, the file will // not be marked as added to the directory. @@ -143,6 +163,24 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) writeFooter(indexPath, engineFileName); } + private int calculateArraySize(float[][] vectors) { + int vectorLength = vectors[0].length; + int numVectors = vectors.length; + int floatByteSize = 4; + int javaReferenceSize = 4; + int javaArrayHeaderSize = 12; + int javaRoundingNumber = 8; + int vectorSize = vectorLength * floatByteSize + javaArrayHeaderSize; + if (vectorSize % javaRoundingNumber != 0) { + vectorSize += vectorSize % javaRoundingNumber; + } + int vectorsSize = numVectors * (vectorSize + javaReferenceSize) + javaArrayHeaderSize; + if (vectorsSize % javaRoundingNumber != 0) { + vectorsSize += vectorsSize % javaRoundingNumber; + } + return vectorsSize; + } + private void createKNNIndexFromTemplate(byte[] model, KNNCodecUtil.Pair pair, KNNEngine knnEngine, String indexPath) { Map parameters = ImmutableMap.of( KNNConstants.INDEX_THREAD_QTY, @@ -204,6 +242,8 @@ private void createKNNIndexFromScratch(FieldInfo fieldInfo, KNNCodecUtil.Pair pa @Override public void merge(MergeState mergeState) { try { + System.out.println("#######"); + System.out.println(mergeState); delegatee.merge(mergeState); assert mergeState != null; assert mergeState.mergeFieldInfos != null; diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java index ce04c9078..00426d534 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java @@ -22,7 +22,13 @@ public enum KNNCounter { SCRIPT_QUERY_ERRORS("script_query_errors"), TRAINING_REQUESTS("training_requests"), TRAINING_ERRORS("training_errors"), - KNN_QUERY_WITH_FILTER_REQUESTS("knn_query_with_filter_requests"); + KNN_QUERY_WITH_FILTER_REQUESTS("knn_query_with_filter_requests"), + REFRESH_CURRENT_OPERATIONS("current_operations"), + REFRESH_CURRENT_DOCS("current_docs"), + REFRESH_CURRENT_SIZE_IN_BYTES("current_size_in_bytes"), + REFRESH_TOTAL_OPERATIONS("total_operations"), + REFRESH_TOTAL_TIME_IN_MILLIS("total_time_in_millis"), + REFRESH_TOTAL_DOCS("total_docs"); private String name; private AtomicLong count; diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java index 66b3f215b..a4f72e980 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java @@ -24,6 +24,7 @@ import java.time.temporal.ChronoUnit; import java.util.HashMap; import java.util.Map; +import java.util.function.Supplier; /** * Class represents all stats the plugin keeps track of @@ -84,6 +85,7 @@ private Map> buildStatsMap() { addEngineStats(builder); addScriptStats(builder); addModelStats(builder); + addGraphStats(builder); return builder.build(); } @@ -169,4 +171,26 @@ private void addModelStats(ImmutableMap.Builder> builder) { new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeAsPercentage)) ); } + + private void addGraphStats(ImmutableMap.Builder> builder) { + builder.put(StatNames.GRAPH_STATS.getName(), new KNNStat<>(false, new Supplier>>() { + @Override + public Map> get() { + return createGraphStatsMap(); + } + })); + } + + private Map> createGraphStatsMap() { + Map refreshMap = new HashMap<>(); + refreshMap.put(KNNCounter.REFRESH_CURRENT_OPERATIONS.getName(), KNNCounter.REFRESH_CURRENT_OPERATIONS.getCount()); + refreshMap.put(KNNCounter.REFRESH_CURRENT_DOCS.getName(), KNNCounter.REFRESH_CURRENT_DOCS.getCount()); + refreshMap.put(KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getName(), KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getCount()); + refreshMap.put(KNNCounter.REFRESH_TOTAL_OPERATIONS.getName(), KNNCounter.REFRESH_TOTAL_OPERATIONS.getCount()); + refreshMap.put(KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.getName(), KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.getCount()); + refreshMap.put(KNNCounter.REFRESH_TOTAL_DOCS.getName(), KNNCounter.REFRESH_TOTAL_DOCS.getCount()); + Map> graphStatsMap = new HashMap<>(); + graphStatsMap.put(StatNames.REFRESH.getName(), refreshMap); + return graphStatsMap; + } } diff --git a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java index a098dd8b5..85b4ca70a 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java @@ -41,7 +41,9 @@ public enum StatNames { TRAINING_MEMORY_USAGE("training_memory_usage"), TRAINING_MEMORY_USAGE_PERCENTAGE("training_memory_usage_percentage"), SCRIPT_QUERY_ERRORS(KNNCounter.SCRIPT_QUERY_ERRORS.getName()), - KNN_QUERY_WITH_FILTER_REQUESTS(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.getName()); + KNN_QUERY_WITH_FILTER_REQUESTS(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.getName()), + GRAPH_STATS("graph_stats"), + REFRESH("refresh"); private String name; From 27acb63bf0bd83c7b33ac119c86821d38fb75e78 Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Thu, 14 Sep 2023 10:27:59 -0700 Subject: [PATCH 02/10] Extract variables to new class and add one test Signed-off-by: Ryan Bogan --- .../KNN80Codec/KNN80DocValuesConsumer.java | 59 ++++++++++------ .../knn/plugin/stats/KNNGraphValue.java | 70 +++++++++++++++++++ .../opensearch/knn/plugin/stats/KNNStats.java | 17 +++-- .../knn/plugin/stats/StatNames.java | 4 +- .../KNN80DocValuesConsumerTests.java | 11 ++- 5 files changed, 132 insertions(+), 29 deletions(-) create mode 100644 src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index 023607de4..a75da3f67 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -36,6 +36,7 @@ import org.apache.lucene.store.FilterDirectory; import org.opensearch.knn.index.mapper.KNNVectorFieldMapper; import org.opensearch.knn.common.KNNConstants; +import org.opensearch.knn.plugin.stats.KNNGraphValue; import java.io.Closeable; import java.io.IOException; @@ -78,10 +79,10 @@ public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) th if (isKNNBinaryFieldRequired(field)) { StopWatch stopWatch = new StopWatch(); stopWatch.start(); - addKNNBinaryField(field, valuesProducer); + addKNNBinaryField(field, valuesProducer, false, true); stopWatch.stop(); long time_in_millis = stopWatch.totalTime().millis(); - KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.set(KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.getCount() + time_in_millis); + KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.set(KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getValue() + time_in_millis); } } @@ -102,7 +103,7 @@ private KNNEngine getKNNEngine(@NonNull FieldInfo field) { return KNNEngine.getEngine(engineName); } - public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, boolean isMerge, boolean isRefresh) throws IOException { // Get values to be index BinaryDocValues values = valuesProducer.getBinary(field); KNNCodecUtil.Pair pair = KNNCodecUtil.getFloats(values); @@ -110,12 +111,14 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) logger.info("Skipping engine index creation as there are no vectors or docs in the documents"); return; } - KNNCounter.REFRESH_CURRENT_OPERATIONS.increment(); - KNNCounter.REFRESH_CURRENT_DOCS.set(KNNCounter.REFRESH_CURRENT_DOCS.getCount() + pair.docs.length); - KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.set( - KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getCount() + calculateArraySize(pair.vectors) - ); - + long arraySize = calculateArraySize(pair.vectors); + if (isMerge) { + KNNGraphValue.MERGE_CURRENT_OPERATIONS.increment(); + KNNGraphValue.MERGE_CURRENT_DOCS.set(KNNGraphValue.MERGE_CURRENT_DOCS.getValue() + pair.docs.length); + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.set( + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() + arraySize + ); + } // Increment counter for number of graph index requests KNNCounter.GRAPH_INDEX_REQUESTS.increment(); // Create library index either from model or from scratch @@ -146,14 +149,22 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) indexCreator = () -> createKNNIndexFromScratch(field, pair, knnEngine, indexPath); } - KNNCounter.REFRESH_CURRENT_OPERATIONS.set(KNNCounter.REFRESH_CURRENT_OPERATIONS.getCount() - 1); - KNNCounter.REFRESH_CURRENT_DOCS.set(KNNCounter.REFRESH_CURRENT_DOCS.getCount() - pair.docs.length); - KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.set( - KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getCount() - calculateArraySize(pair.vectors) - ); - KNNCounter.REFRESH_TOTAL_OPERATIONS.increment(); - System.out.println(KNNCounter.REFRESH_TOTAL_OPERATIONS.getCount()); - KNNCounter.REFRESH_TOTAL_DOCS.set(KNNCounter.REFRESH_TOTAL_DOCS.getCount() + pair.docs.length); + if (isMerge) { + KNNGraphValue.MERGE_CURRENT_OPERATIONS.set(KNNGraphValue.MERGE_CURRENT_OPERATIONS.getValue() - 1); + KNNGraphValue.MERGE_CURRENT_DOCS.set(KNNGraphValue.MERGE_CURRENT_DOCS.getValue() - pair.docs.length); + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.set( + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() - calculateArraySize(pair.vectors) + ); + KNNGraphValue.MERGE_TOTAL_OPERATIONS.increment(); + KNNGraphValue.MERGE_TOTAL_DOCS.set(KNNGraphValue.MERGE_TOTAL_DOCS.getValue() + pair.docs.length); + KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.set( + KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue() + arraySize + ); + } + + if (isRefresh) { + KNNGraphValue.REFRESH_TOTAL_OPERATIONS.increment(); + } // This is a bit of a hack. We have to create an output here and then immediately close it to ensure that // engineFileName is added to the tracked files by Lucene's TrackingDirectoryWrapper. Otherwise, the file will @@ -242,15 +253,23 @@ private void createKNNIndexFromScratch(FieldInfo fieldInfo, KNNCodecUtil.Pair pa @Override public void merge(MergeState mergeState) { try { - System.out.println("#######"); - System.out.println(mergeState); + System.out.println("!!!!!!!#.25"); delegatee.merge(mergeState); + System.out.println("!!!!!!!#.5"); assert mergeState != null; assert mergeState.mergeFieldInfos != null; for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) { + System.out.println("!!!!!!!!#1"); DocValuesType type = fieldInfo.getDocValuesType(); + System.out.println("!!!!!!!!#2"); if (type == DocValuesType.BINARY && fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) { - addKNNBinaryField(fieldInfo, new KNN80DocValuesReader(mergeState)); + System.out.println("!!!!!!!!#3"); + StopWatch stopWatch = new StopWatch(); + stopWatch.start(); + addKNNBinaryField(fieldInfo, new KNN80DocValuesReader(mergeState), true, false); + stopWatch.stop(); + long time_in_millis = stopWatch.totalTime().millis(); + KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.set(KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.getValue() + time_in_millis); } } } catch (Exception e) { diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java new file mode 100644 index 000000000..3588426a8 --- /dev/null +++ b/src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java @@ -0,0 +1,70 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.knn.plugin.stats; + +import java.util.concurrent.atomic.AtomicLong; + +/** + * Contains a map of counters to keep track of different values + */ +public enum KNNGraphValue { + + REFRESH_TOTAL_OPERATIONS("total"), + REFRESH_TOTAL_TIME_IN_MILLIS("total_time_in_millis"), + MERGE_CURRENT_OPERATIONS("current"), + MERGE_CURRENT_DOCS("current_docs"), + MERGE_CURRENT_SIZE_IN_BYTES("current_size_in_bytes"), + MERGE_TOTAL_OPERATIONS("total"), + MERGE_TOTAL_TIME_IN_MILLIS("total_time_in_millis"), + MERGE_TOTAL_DOCS("total_docs"), + MERGE_TOTAL_SIZE_IN_BYTES("total_size_in_bytes"); + + private String name; + private AtomicLong value; + + /** + * Constructor + * + * @param name name of the counter + */ + KNNGraphValue(String name) { + this.name = name; + this.value = new AtomicLong(0); + } + + /** + * Get name of value + * + * @return name + */ + public String getName() { + return name; + } + + /** + * Get the value of count + * + * @return count + */ + public Long getValue() { + return value.get(); + } + + /** + * Increment the value of a counter + */ + public void increment() { + value.getAndIncrement(); + } + + /** + * @param value counter value + * Set the value of a counter + */ + public void set(long value) { + this.value.set(value); + } +} diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java index a4f72e980..07d129652 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/KNNStats.java @@ -182,14 +182,19 @@ public Map> get() { } private Map> createGraphStatsMap() { + Map mergeMap = new HashMap<>(); + mergeMap.put(KNNGraphValue.MERGE_CURRENT_OPERATIONS.getName(), KNNGraphValue.MERGE_CURRENT_OPERATIONS.getValue()); + mergeMap.put(KNNGraphValue.MERGE_CURRENT_DOCS.getName(), KNNGraphValue.MERGE_CURRENT_DOCS.getValue()); + mergeMap.put(KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getName(), KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue()); + mergeMap.put(KNNGraphValue.MERGE_TOTAL_OPERATIONS.getName(), KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue()); + mergeMap.put(KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.getName(), KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.getValue()); + mergeMap.put(KNNGraphValue.MERGE_TOTAL_DOCS.getName(), KNNGraphValue.MERGE_TOTAL_DOCS.getValue()); + mergeMap.put(KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getName(), KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue()); Map refreshMap = new HashMap<>(); - refreshMap.put(KNNCounter.REFRESH_CURRENT_OPERATIONS.getName(), KNNCounter.REFRESH_CURRENT_OPERATIONS.getCount()); - refreshMap.put(KNNCounter.REFRESH_CURRENT_DOCS.getName(), KNNCounter.REFRESH_CURRENT_DOCS.getCount()); - refreshMap.put(KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getName(), KNNCounter.REFRESH_CURRENT_SIZE_IN_BYTES.getCount()); - refreshMap.put(KNNCounter.REFRESH_TOTAL_OPERATIONS.getName(), KNNCounter.REFRESH_TOTAL_OPERATIONS.getCount()); - refreshMap.put(KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.getName(), KNNCounter.REFRESH_TOTAL_TIME_IN_MILLIS.getCount()); - refreshMap.put(KNNCounter.REFRESH_TOTAL_DOCS.getName(), KNNCounter.REFRESH_TOTAL_DOCS.getCount()); + refreshMap.put(KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getName(), KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + refreshMap.put(KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getName(), KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getValue()); Map> graphStatsMap = new HashMap<>(); + graphStatsMap.put(StatNames.MERGE.getName(), mergeMap); graphStatsMap.put(StatNames.REFRESH.getName(), refreshMap); return graphStatsMap; } diff --git a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java index 85b4ca70a..11456dc3f 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java @@ -43,7 +43,9 @@ public enum StatNames { SCRIPT_QUERY_ERRORS(KNNCounter.SCRIPT_QUERY_ERRORS.getName()), KNN_QUERY_WITH_FILTER_REQUESTS(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.getName()), GRAPH_STATS("graph_stats"), - REFRESH("refresh"); + REFRESH("refresh"), + + MERGE("merge"); private String name; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java index 58f4b6e39..dd347cece 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java @@ -39,9 +39,11 @@ import org.opensearch.knn.indices.ModelState; import org.opensearch.knn.jni.JNIService; import org.opensearch.knn.plugin.stats.KNNCounter; +import org.opensearch.knn.plugin.stats.KNNGraphValue; import java.io.IOException; import java.util.Map; +import java.util.Optional; import java.util.concurrent.ExecutionException; import static org.mockito.ArgumentMatchers.any; @@ -348,7 +350,7 @@ public void testAddKNNBinaryField_fromModel_faiss() throws IOException, Executio // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, true); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -359,6 +361,11 @@ public void testAddKNNBinaryField_fromModel_faiss() throws IOException, Executio // The document should be readable by faiss assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + + // The refresh statistics should be updated + assertEquals(1, (long) KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getValue()); + } public void testMerge_exception() throws IOException { @@ -426,6 +433,6 @@ public void testAddBinaryField_luceneEngine_noInvocations_addKNNBinary() throws knn80DocValuesConsumer.addBinaryField(fieldInfo, docValuesProducer); verify(delegate, times(1)).addBinaryField(fieldInfo, docValuesProducer); - verify(knn80DocValuesConsumer, never()).addKNNBinaryField(any(), any()); + verify(knn80DocValuesConsumer, never()).addKNNBinaryField(any(), any(), false, false); } } From 104da3b51fb9ca9bc1adedc3c207ed7f0dfdd369 Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Mon, 18 Sep 2023 09:44:28 -0700 Subject: [PATCH 03/10] Minor changes and add tests Signed-off-by: Ryan Bogan --- .../KNN80Codec/KNN80DocValuesConsumer.java | 19 ++++++------------- .../knn/plugin/stats/KNNCounter.java | 8 +------- .../KNN80DocValuesConsumerTests.java | 17 ++++++++--------- 3 files changed, 15 insertions(+), 29 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index a75da3f67..770249733 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -76,6 +76,7 @@ class KNN80DocValuesConsumer extends DocValuesConsumer implements Closeable { @Override public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException { + delegatee.addBinaryField(field, valuesProducer); if (isKNNBinaryFieldRequired(field)) { StopWatch stopWatch = new StopWatch(); stopWatch.start(); @@ -103,7 +104,8 @@ private KNNEngine getKNNEngine(@NonNull FieldInfo field) { return KNNEngine.getEngine(engineName); } - public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, boolean isMerge, boolean isRefresh) throws IOException { + public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, boolean isMerge, boolean isRefresh) + throws IOException { // Get values to be index BinaryDocValues values = valuesProducer.getBinary(field); KNNCodecUtil.Pair pair = KNNCodecUtil.getFloats(values); @@ -115,9 +117,7 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, if (isMerge) { KNNGraphValue.MERGE_CURRENT_OPERATIONS.increment(); KNNGraphValue.MERGE_CURRENT_DOCS.set(KNNGraphValue.MERGE_CURRENT_DOCS.getValue() + pair.docs.length); - KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.set( - KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() + arraySize - ); + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.set(KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() + arraySize); } // Increment counter for number of graph index requests KNNCounter.GRAPH_INDEX_REQUESTS.increment(); @@ -153,13 +153,11 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, KNNGraphValue.MERGE_CURRENT_OPERATIONS.set(KNNGraphValue.MERGE_CURRENT_OPERATIONS.getValue() - 1); KNNGraphValue.MERGE_CURRENT_DOCS.set(KNNGraphValue.MERGE_CURRENT_DOCS.getValue() - pair.docs.length); KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.set( - KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() - calculateArraySize(pair.vectors) + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() - calculateArraySize(pair.vectors) ); KNNGraphValue.MERGE_TOTAL_OPERATIONS.increment(); KNNGraphValue.MERGE_TOTAL_DOCS.set(KNNGraphValue.MERGE_TOTAL_DOCS.getValue() + pair.docs.length); - KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.set( - KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue() + arraySize - ); + KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.set(KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue() + arraySize); } if (isRefresh) { @@ -253,17 +251,12 @@ private void createKNNIndexFromScratch(FieldInfo fieldInfo, KNNCodecUtil.Pair pa @Override public void merge(MergeState mergeState) { try { - System.out.println("!!!!!!!#.25"); delegatee.merge(mergeState); - System.out.println("!!!!!!!#.5"); assert mergeState != null; assert mergeState.mergeFieldInfos != null; for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) { - System.out.println("!!!!!!!!#1"); DocValuesType type = fieldInfo.getDocValuesType(); - System.out.println("!!!!!!!!#2"); if (type == DocValuesType.BINARY && fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) { - System.out.println("!!!!!!!!#3"); StopWatch stopWatch = new StopWatch(); stopWatch.start(); addKNNBinaryField(fieldInfo, new KNN80DocValuesReader(mergeState), true, false); diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java index 00426d534..ce04c9078 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/KNNCounter.java @@ -22,13 +22,7 @@ public enum KNNCounter { SCRIPT_QUERY_ERRORS("script_query_errors"), TRAINING_REQUESTS("training_requests"), TRAINING_ERRORS("training_errors"), - KNN_QUERY_WITH_FILTER_REQUESTS("knn_query_with_filter_requests"), - REFRESH_CURRENT_OPERATIONS("current_operations"), - REFRESH_CURRENT_DOCS("current_docs"), - REFRESH_CURRENT_SIZE_IN_BYTES("current_size_in_bytes"), - REFRESH_TOTAL_OPERATIONS("total_operations"), - REFRESH_TOTAL_TIME_IN_MILLIS("total_time_in_millis"), - REFRESH_TOTAL_DOCS("total_docs"); + KNN_QUERY_WITH_FILTER_REQUESTS("knn_query_with_filter_requests"); private String name; private AtomicLong count; diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java index dd347cece..5c30935d5 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java @@ -43,10 +43,10 @@ import java.io.IOException; import java.util.Map; -import java.util.Optional; import java.util.concurrent.ExecutionException; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -96,7 +96,7 @@ public void testAddBinaryField_withKNN() throws IOException { KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(delegate, null) { @Override - public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) { + public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, boolean isMerge, boolean isRefresh) { called[0] = true; } }; @@ -120,7 +120,7 @@ public void testAddBinaryField_withoutKNN() throws IOException { KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(delegate, null) { @Override - public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer) { + public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, boolean isMerge, boolean isRefresh) { called[0] = true; } }; @@ -136,7 +136,7 @@ public void testAddKNNBinaryField_noVectors() throws IOException { RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(0, 128); Long initialGraphIndexRequests = KNNCounter.GRAPH_INDEX_REQUESTS.getCount(); KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, null); - knn80DocValuesConsumer.addKNNBinaryField(null, randomVectorDocValuesProducer); + knn80DocValuesConsumer.addKNNBinaryField(null, randomVectorDocValuesProducer, false, false); assertEquals(initialGraphIndexRequests, KNNCounter.GRAPH_INDEX_REQUESTS.getCount()); } @@ -179,7 +179,7 @@ public void testAddKNNBinaryField_fromScratch_nmslibCurrent() throws IOException // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, false); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -223,7 +223,7 @@ public void testAddKNNBinaryField_fromScratch_nmslibLegacy() throws IOException // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, false); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -274,7 +274,7 @@ public void testAddKNNBinaryField_fromScratch_faissCurrent() throws IOException // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, false); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -364,7 +364,6 @@ public void testAddKNNBinaryField_fromModel_faiss() throws IOException, Executio // The refresh statistics should be updated assertEquals(1, (long) KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); - assertNotEquals(0, (long) KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getValue()); } @@ -433,6 +432,6 @@ public void testAddBinaryField_luceneEngine_noInvocations_addKNNBinary() throws knn80DocValuesConsumer.addBinaryField(fieldInfo, docValuesProducer); verify(delegate, times(1)).addBinaryField(fieldInfo, docValuesProducer); - verify(knn80DocValuesConsumer, never()).addKNNBinaryField(any(), any(), false, false); + verify(knn80DocValuesConsumer, never()).addKNNBinaryField(any(), any(), eq(false), eq(true)); } } From ed14e8bd43cc1a426ee07e122d18d4ae1b78f7b1 Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Mon, 18 Sep 2023 10:35:34 -0700 Subject: [PATCH 04/10] Change tests Signed-off-by: Ryan Bogan --- .../KNN80DocValuesConsumerTests.java | 55 ++++++++++++++++--- 1 file changed, 48 insertions(+), 7 deletions(-) diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java index 5c30935d5..6af83de87 100644 --- a/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java +++ b/src/test/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumerTests.java @@ -135,9 +135,17 @@ public void testAddKNNBinaryField_noVectors() throws IOException { // When there are no new vectors, no more graph index requests should be added RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(0, 128); Long initialGraphIndexRequests = KNNCounter.GRAPH_INDEX_REQUESTS.getCount(); + Long initialRefreshOperations = KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue(); + Long initialMergeOperations = KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue(); + Long initialMergeSize = KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue(); + Long initialMergeDocs = KNNGraphValue.MERGE_TOTAL_DOCS.getValue(); KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, null); - knn80DocValuesConsumer.addKNNBinaryField(null, randomVectorDocValuesProducer, false, false); + knn80DocValuesConsumer.addKNNBinaryField(null, randomVectorDocValuesProducer, true, true); assertEquals(initialGraphIndexRequests, KNNCounter.GRAPH_INDEX_REQUESTS.getCount()); + assertEquals(initialRefreshOperations, KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + assertEquals(initialMergeOperations, KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue()); + assertEquals(initialMergeSize, KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue()); + assertEquals(initialMergeDocs, KNNGraphValue.MERGE_TOTAL_DOCS.getValue()); } public void testAddKNNBinaryField_fromScratch_nmslibCurrent() throws IOException { @@ -176,10 +184,13 @@ public void testAddKNNBinaryField_fromScratch_nmslibCurrent() throws IOException FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + long initialRefreshOperations = KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue(); + long initialMergeOperations = KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue(); + // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, false); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, true, true); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -190,6 +201,12 @@ public void testAddKNNBinaryField_fromScratch_nmslibCurrent() throws IOException // The document should be readable by nmslib assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + + // The graph creation statistics should be updated + assertEquals(1 + initialRefreshOperations, (long) KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + assertEquals(1 + initialMergeOperations, (long) KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_DOCS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue()); } public void testAddKNNBinaryField_fromScratch_nmslibLegacy() throws IOException { @@ -220,10 +237,13 @@ public void testAddKNNBinaryField_fromScratch_nmslibLegacy() throws IOException FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + long initialRefreshOperations = KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue(); + long initialMergeOperations = KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue(); + // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, false); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, true, true); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -234,6 +254,12 @@ public void testAddKNNBinaryField_fromScratch_nmslibLegacy() throws IOException // The document should be readable by nmslib assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + + // The graph creation statistics should be updated + assertEquals(1 + initialRefreshOperations, (long) KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + assertEquals(1 + initialMergeOperations, (long) KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_DOCS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue()); } public void testAddKNNBinaryField_fromScratch_faissCurrent() throws IOException { @@ -271,10 +297,13 @@ public void testAddKNNBinaryField_fromScratch_faissCurrent() throws IOException FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + long initialRefreshOperations = KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue(); + long initialMergeOperations = KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue(); + // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, false); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, true, true); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -285,6 +314,12 @@ public void testAddKNNBinaryField_fromScratch_faissCurrent() throws IOException // The document should be readable by faiss assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); + + // The graph creation statistics should be updated + assertEquals(1 + initialRefreshOperations, (long) KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + assertEquals(1 + initialMergeOperations, (long) KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_DOCS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue()); } public void testAddKNNBinaryField_fromModel_faiss() throws IOException, ExecutionException, InterruptedException { @@ -347,10 +382,13 @@ public void testAddKNNBinaryField_fromModel_faiss() throws IOException, Executio FieldInfos fieldInfos = new FieldInfos(fieldInfoArray); SegmentWriteState state = new SegmentWriteState(null, directory, segmentInfo, fieldInfos, null, IOContext.DEFAULT); + long initialRefreshOperations = KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue(); + long initialMergeOperations = KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue(); + // Add documents to the field KNN80DocValuesConsumer knn80DocValuesConsumer = new KNN80DocValuesConsumer(null, state); RandomVectorDocValuesProducer randomVectorDocValuesProducer = new RandomVectorDocValuesProducer(docsInSegment, dimension); - knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, false, true); + knn80DocValuesConsumer.addKNNBinaryField(fieldInfoArray[0], randomVectorDocValuesProducer, true, true); // The document should be created in the correct location String expectedFile = KNNCodecUtil.buildEngineFileName(segmentName, knnEngine.getVersion(), fieldName, knnEngine.getExtension()); @@ -362,8 +400,11 @@ public void testAddKNNBinaryField_fromModel_faiss() throws IOException, Executio // The document should be readable by faiss assertLoadableByEngine(state, expectedFile, knnEngine, spaceType, dimension); - // The refresh statistics should be updated - assertEquals(1, (long) KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + // The graph creation statistics should be updated + assertEquals(1 + initialRefreshOperations, (long) KNNGraphValue.REFRESH_TOTAL_OPERATIONS.getValue()); + assertEquals(1 + initialMergeOperations, (long) KNNGraphValue.MERGE_TOTAL_OPERATIONS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_DOCS.getValue()); + assertNotEquals(0, (long) KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue()); } From 3394066148ae1b25aa122855ee2736ea8944016f Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Mon, 18 Sep 2023 10:38:30 -0700 Subject: [PATCH 05/10] Update CHANGELOG Signed-off-by: Ryan Bogan --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index f0832aa6e..3dbfffa86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.10...2.x) ### Features ### Enhancements +- Add graph creation stats to the KNNStats API. [#1141](https://github.com/opensearch-project/k-NN/pull/1141) ### Bug Fixes ### Infrastructure ### Documentation From 82df6c96aa6e55ba2230d877a73d7672ef49a0cc Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Mon, 18 Sep 2023 10:42:28 -0700 Subject: [PATCH 06/10] Fix typos Signed-off-by: Ryan Bogan --- .../knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java | 2 +- src/main/java/org/opensearch/knn/plugin/stats/StatNames.java | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index 770249733..2a085a789 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -106,7 +106,7 @@ private KNNEngine getKNNEngine(@NonNull FieldInfo field) { public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, boolean isMerge, boolean isRefresh) throws IOException { - // Get values to be index + // Get values to be indexed BinaryDocValues values = valuesProducer.getBinary(field); KNNCodecUtil.Pair pair = KNNCodecUtil.getFloats(values); if (pair.vectors.length == 0 || pair.docs.length == 0) { diff --git a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java index 11456dc3f..e9ed2b126 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/StatNames.java @@ -44,7 +44,6 @@ public enum StatNames { KNN_QUERY_WITH_FILTER_REQUESTS(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.getName()), GRAPH_STATS("graph_stats"), REFRESH("refresh"), - MERGE("merge"); private String name; From cd41e46f90d09ffa9ae7bd552c3bc9f323209e51 Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Fri, 22 Sep 2023 09:21:26 -0700 Subject: [PATCH 07/10] Address PR feedback Signed-off-by: Ryan Bogan --- .../KNN80Codec/KNN80DocValuesConsumer.java | 43 +++++++------------ .../knn/index/codec/util/KNNCodecUtil.java | 22 ++++++++++ .../knn/plugin/stats/KNNGraphValue.java | 39 ++++++++++++++--- 3 files changed, 70 insertions(+), 34 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index 2a085a789..1b90b5db6 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -55,6 +55,7 @@ import static org.opensearch.knn.common.KNNConstants.MODEL_ID; import static org.opensearch.knn.common.KNNConstants.PARAMETERS; import static org.opensearch.knn.index.codec.util.KNNCodecUtil.buildEngineFileName; +import static org.opensearch.knn.index.codec.util.KNNCodecUtil.calculateArraySize; /** * This class writes the KNN docvalues to the segments @@ -116,8 +117,8 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, long arraySize = calculateArraySize(pair.vectors); if (isMerge) { KNNGraphValue.MERGE_CURRENT_OPERATIONS.increment(); - KNNGraphValue.MERGE_CURRENT_DOCS.set(KNNGraphValue.MERGE_CURRENT_DOCS.getValue() + pair.docs.length); - KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.set(KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() + arraySize); + KNNGraphValue.MERGE_CURRENT_DOCS.incrementBy(pair.docs.length); + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.incrementBy(arraySize); } // Increment counter for number of graph index requests KNNCounter.GRAPH_INDEX_REQUESTS.increment(); @@ -150,18 +151,11 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, } if (isMerge) { - KNNGraphValue.MERGE_CURRENT_OPERATIONS.set(KNNGraphValue.MERGE_CURRENT_OPERATIONS.getValue() - 1); - KNNGraphValue.MERGE_CURRENT_DOCS.set(KNNGraphValue.MERGE_CURRENT_DOCS.getValue() - pair.docs.length); - KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.set( - KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.getValue() - calculateArraySize(pair.vectors) - ); - KNNGraphValue.MERGE_TOTAL_OPERATIONS.increment(); - KNNGraphValue.MERGE_TOTAL_DOCS.set(KNNGraphValue.MERGE_TOTAL_DOCS.getValue() + pair.docs.length); - KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.set(KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.getValue() + arraySize); + recordMergeStats(pair.docs.length, arraySize); } if (isRefresh) { - KNNGraphValue.REFRESH_TOTAL_OPERATIONS.increment(); + recordRefreshStats(); } // This is a bit of a hack. We have to create an output here and then immediately close it to ensure that @@ -172,22 +166,17 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, writeFooter(indexPath, engineFileName); } - private int calculateArraySize(float[][] vectors) { - int vectorLength = vectors[0].length; - int numVectors = vectors.length; - int floatByteSize = 4; - int javaReferenceSize = 4; - int javaArrayHeaderSize = 12; - int javaRoundingNumber = 8; - int vectorSize = vectorLength * floatByteSize + javaArrayHeaderSize; - if (vectorSize % javaRoundingNumber != 0) { - vectorSize += vectorSize % javaRoundingNumber; - } - int vectorsSize = numVectors * (vectorSize + javaReferenceSize) + javaArrayHeaderSize; - if (vectorsSize % javaRoundingNumber != 0) { - vectorsSize += vectorsSize % javaRoundingNumber; - } - return vectorsSize; + private void recordMergeStats(int length, long arraySize) { + KNNGraphValue.MERGE_CURRENT_OPERATIONS.decrement(); + KNNGraphValue.MERGE_CURRENT_DOCS.decrementBy(length); + KNNGraphValue.MERGE_CURRENT_SIZE_IN_BYTES.decrementBy(arraySize); + KNNGraphValue.MERGE_TOTAL_OPERATIONS.increment(); + KNNGraphValue.MERGE_TOTAL_DOCS.incrementBy(length); + KNNGraphValue.MERGE_TOTAL_SIZE_IN_BYTES.incrementBy(arraySize); + } + + private void recordRefreshStats() { + KNNGraphValue.REFRESH_TOTAL_OPERATIONS.increment(); } private void createKNNIndexFromTemplate(byte[] model, KNNCodecUtil.Pair pair, KNNEngine knnEngine, String indexPath) { diff --git a/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java b/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java index eef9e6863..c07bcc4eb 100644 --- a/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java +++ b/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java @@ -17,6 +17,14 @@ public class KNNCodecUtil { public static final String HNSW_EXTENSION = ".hnsw"; public static final String HNSW_COMPOUND_EXTENSION = ".hnswc"; + // Floats are 4 bytes in size + public static final int FLOAT_BYTE_SIZE = 4; + // References to objects are 4 bytes in size + public static final int JAVA_REFERENCE_SIZE = 4; + // Each array in Java has a header that is 12 bytes + public static final int JAVA_ARRAY_HEADER_SIZE = 12; + // Java rounds each array size up to multiples of 8 bytes + public static final int JAVA_ROUNDING_NUMBER = 8; public static final class Pair { public Pair(int[] docs, float[][] vectors) { @@ -43,6 +51,20 @@ public static KNNCodecUtil.Pair getFloats(BinaryDocValues values) throws IOExcep return new KNNCodecUtil.Pair(docIdList.stream().mapToInt(Integer::intValue).toArray(), vectorList.toArray(new float[][] {})); } + public static long calculateArraySize(float[][] vectors) { + int vectorLength = vectors[0].length; + int numVectors = vectors.length; + int vectorSize = vectorLength * FLOAT_BYTE_SIZE + JAVA_ARRAY_HEADER_SIZE; + if (vectorSize % JAVA_ROUNDING_NUMBER != 0) { + vectorSize += vectorSize % JAVA_ROUNDING_NUMBER; + } + int vectorsSize = numVectors * (vectorSize + JAVA_REFERENCE_SIZE) + JAVA_ARRAY_HEADER_SIZE; + if (vectorsSize % JAVA_ROUNDING_NUMBER != 0) { + vectorsSize += vectorsSize % JAVA_ROUNDING_NUMBER; + } + return vectorsSize; + } + public static String buildEngineFileName(String segmentName, String latestBuildVersion, String fieldName, String extension) { return String.format("%s%s%s", buildEngineFilePrefix(segmentName), latestBuildVersion, buildEngineFileSuffix(fieldName, extension)); } diff --git a/src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java b/src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java index 3588426a8..b33b59e36 100644 --- a/src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java +++ b/src/main/java/org/opensearch/knn/plugin/stats/KNNGraphValue.java @@ -8,7 +8,7 @@ import java.util.concurrent.atomic.AtomicLong; /** - * Contains a map of counters to keep track of different values + * Contains a map to keep track of different graph values */ public enum KNNGraphValue { @@ -28,7 +28,7 @@ public enum KNNGraphValue { /** * Constructor * - * @param name name of the counter + * @param name name of the graph value */ KNNGraphValue(String name) { this.name = name; @@ -45,24 +45,49 @@ public String getName() { } /** - * Get the value of count + * Get the graph value * - * @return count + * @return value */ public Long getValue() { return value.get(); } /** - * Increment the value of a counter + * Increment the graph value */ public void increment() { value.getAndIncrement(); } /** - * @param value counter value - * Set the value of a counter + * Decrement the graph value + */ + public void decrement() { + value.getAndDecrement(); + } + + /** + * Increment the graph value by a specified amount + * + * @param delta The amount to increment + */ + public void incrementBy(long delta) { + value.getAndAdd(delta); + } + + /** + * Decrement the graph value by a specified amount + * + * @param delta The amount to decrement + */ + public void decrementBy(long delta) { + value.set(value.get() - delta); + } + + /** + * @param value graph value + * Set the graph value */ public void set(long value) { this.value.set(value); From f3e6722435fbb8a073cc54a331a9596e1c69f1cb Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Mon, 25 Sep 2023 10:15:26 -0700 Subject: [PATCH 08/10] Add logging for time taken by merge/refresh operations Signed-off-by: Ryan Bogan --- .../knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index 1b90b5db6..cdfddf1d9 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -85,6 +85,7 @@ public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) th stopWatch.stop(); long time_in_millis = stopWatch.totalTime().millis(); KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.set(KNNGraphValue.REFRESH_TOTAL_TIME_IN_MILLIS.getValue() + time_in_millis); + logger.warn("Refresh operation complete in " + time_in_millis + " ms"); } } @@ -252,6 +253,7 @@ public void merge(MergeState mergeState) { stopWatch.stop(); long time_in_millis = stopWatch.totalTime().millis(); KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.set(KNNGraphValue.MERGE_TOTAL_TIME_IN_MILLIS.getValue() + time_in_millis); + logger.warn("Merge operation complete in " + time_in_millis + " ms"); } } } catch (Exception e) { From fe891e6b9a4c85314e910b82c03f6b6ecbe969a6 Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Wed, 27 Sep 2023 10:42:07 -0700 Subject: [PATCH 09/10] Change array calculations Signed-off-by: Ryan Bogan --- .../KNN80Codec/KNN80DocValuesConsumer.java | 2 +- .../knn/index/codec/util/KNNCodecUtil.java | 42 ++++++++++++++----- .../util/KNNVectorSerializerFactory.java | 2 +- 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java index cdfddf1d9..bcc2bf369 100644 --- a/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java +++ b/src/main/java/org/opensearch/knn/index/codec/KNN80Codec/KNN80DocValuesConsumer.java @@ -115,7 +115,7 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer, logger.info("Skipping engine index creation as there are no vectors or docs in the documents"); return; } - long arraySize = calculateArraySize(pair.vectors); + long arraySize = calculateArraySize(pair.vectors, pair.serializationMode); if (isMerge) { KNNGraphValue.MERGE_CURRENT_OPERATIONS.increment(); KNNGraphValue.MERGE_CURRENT_DOCS.incrementBy(pair.docs.length); diff --git a/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java b/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java index c07bcc4eb..02ab2d833 100644 --- a/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java +++ b/src/main/java/org/opensearch/knn/index/codec/util/KNNCodecUtil.java @@ -27,42 +27,62 @@ public class KNNCodecUtil { public static final int JAVA_ROUNDING_NUMBER = 8; public static final class Pair { - public Pair(int[] docs, float[][] vectors) { + public Pair(int[] docs, float[][] vectors, SerializationMode serializationMode) { this.docs = docs; this.vectors = vectors; + this.serializationMode = serializationMode; } public int[] docs; public float[][] vectors; + public SerializationMode serializationMode; } public static KNNCodecUtil.Pair getFloats(BinaryDocValues values) throws IOException { ArrayList vectorList = new ArrayList<>(); ArrayList docIdList = new ArrayList<>(); + SerializationMode serializationMode = SerializationMode.COLLECTION_OF_FLOATS; for (int doc = values.nextDoc(); doc != DocIdSetIterator.NO_MORE_DOCS; doc = values.nextDoc()) { BytesRef bytesref = values.binaryValue(); try (ByteArrayInputStream byteStream = new ByteArrayInputStream(bytesref.bytes, bytesref.offset, bytesref.length)) { + serializationMode = KNNVectorSerializerFactory.serializerModeFromStream(byteStream); final KNNVectorSerializer vectorSerializer = KNNVectorSerializerFactory.getSerializerByStreamContent(byteStream); final float[] vector = vectorSerializer.byteToFloatArray(byteStream); vectorList.add(vector); } docIdList.add(doc); } - return new KNNCodecUtil.Pair(docIdList.stream().mapToInt(Integer::intValue).toArray(), vectorList.toArray(new float[][] {})); + return new KNNCodecUtil.Pair( + docIdList.stream().mapToInt(Integer::intValue).toArray(), + vectorList.toArray(new float[][] {}), + serializationMode + ); } - public static long calculateArraySize(float[][] vectors) { + public static long calculateArraySize(float[][] vectors, SerializationMode serializationMode) { int vectorLength = vectors[0].length; int numVectors = vectors.length; - int vectorSize = vectorLength * FLOAT_BYTE_SIZE + JAVA_ARRAY_HEADER_SIZE; - if (vectorSize % JAVA_ROUNDING_NUMBER != 0) { - vectorSize += vectorSize % JAVA_ROUNDING_NUMBER; - } - int vectorsSize = numVectors * (vectorSize + JAVA_REFERENCE_SIZE) + JAVA_ARRAY_HEADER_SIZE; - if (vectorsSize % JAVA_ROUNDING_NUMBER != 0) { - vectorsSize += vectorsSize % JAVA_ROUNDING_NUMBER; + if (serializationMode == SerializationMode.ARRAY) { + int vectorSize = vectorLength * FLOAT_BYTE_SIZE + JAVA_ARRAY_HEADER_SIZE; + if (vectorSize % JAVA_ROUNDING_NUMBER != 0) { + vectorSize += vectorSize % JAVA_ROUNDING_NUMBER; + } + int vectorsSize = numVectors * (vectorSize + JAVA_REFERENCE_SIZE) + JAVA_ARRAY_HEADER_SIZE; + if (vectorsSize % JAVA_ROUNDING_NUMBER != 0) { + vectorsSize += vectorsSize % JAVA_ROUNDING_NUMBER; + } + return vectorsSize; + } else { + int vectorSize = vectorLength * FLOAT_BYTE_SIZE; + if (vectorSize % JAVA_ROUNDING_NUMBER != 0) { + vectorSize += vectorSize % JAVA_ROUNDING_NUMBER; + } + int vectorsSize = numVectors * (vectorSize + JAVA_REFERENCE_SIZE); + if (vectorsSize % JAVA_ROUNDING_NUMBER != 0) { + vectorsSize += vectorsSize % JAVA_ROUNDING_NUMBER; + } + return vectorsSize; } - return vectorsSize; } public static String buildEngineFileName(String segmentName, String latestBuildVersion, String fieldName, String extension) { diff --git a/src/main/java/org/opensearch/knn/index/codec/util/KNNVectorSerializerFactory.java b/src/main/java/org/opensearch/knn/index/codec/util/KNNVectorSerializerFactory.java index f02da0949..5c1e4ca9b 100644 --- a/src/main/java/org/opensearch/knn/index/codec/util/KNNVectorSerializerFactory.java +++ b/src/main/java/org/opensearch/knn/index/codec/util/KNNVectorSerializerFactory.java @@ -56,7 +56,7 @@ public static KNNVectorSerializer getSerializerByStreamContent(final ByteArrayIn return getSerializerBySerializationMode(serializationMode); } - private static SerializationMode serializerModeFromStream(ByteArrayInputStream byteStream) { + static SerializationMode serializerModeFromStream(ByteArrayInputStream byteStream) { int numberOfAvailableBytesInStream = byteStream.available(); if (numberOfAvailableBytesInStream < ARRAY_HEADER_OFFSET) { return getSerializerOrThrowError(numberOfAvailableBytesInStream, COLLECTION_OF_FLOATS); From 4faab0eff6d61d9989d75dcb3728fa7b0f50b73d Mon Sep 17 00:00:00 2001 From: Ryan Bogan Date: Wed, 27 Sep 2023 10:56:09 -0700 Subject: [PATCH 10/10] Minor change Signed-off-by: Ryan Bogan --- src/testFixtures/java/org/opensearch/knn/TestUtils.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/testFixtures/java/org/opensearch/knn/TestUtils.java b/src/testFixtures/java/org/opensearch/knn/TestUtils.java index 12751dff8..f22e4b267 100644 --- a/src/testFixtures/java/org/opensearch/knn/TestUtils.java +++ b/src/testFixtures/java/org/opensearch/knn/TestUtils.java @@ -16,6 +16,7 @@ import java.io.FileReader; import java.io.IOException; import org.opensearch.knn.index.SpaceType; +import org.opensearch.knn.index.codec.util.SerializationMode; import org.opensearch.knn.plugin.script.KNNScoringUtil; import java.util.Comparator; import java.util.Random; @@ -283,7 +284,7 @@ private KNNCodecUtil.Pair readIndexData(String path) throws IOException { } } - return new KNNCodecUtil.Pair(idsArray, vectorsArray); + return new KNNCodecUtil.Pair(idsArray, vectorsArray, SerializationMode.COLLECTION_OF_FLOATS); } private float[][] readQueries(String path) throws IOException {