Add graph creation stats to the KNNStats API

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

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;
@@ -35,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;
@@ -76,7 +78,12 @@ class KNN80DocValuesConsumer extends DocValuesConsumer implements Closeable {
     public void addBinaryField(FieldInfo field, DocValuesProducer valuesProducer) throws IOException {
         delegatee.addBinaryField(field, valuesProducer);
         if (isKNNBinaryFieldRequired(field)) {
-            addKNNBinaryField(field, valuesProducer);
+            StopWatch stopWatch = new StopWatch();
+            stopWatch.start();
+            addKNNBinaryField(field, valuesProducer, false, true);
+            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);
         }
     }
 
@@ -97,14 +104,21 @@ 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 indexed
         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;
         }
+        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
@@ -135,6 +149,21 @@ public void addKNNBinaryField(FieldInfo field, DocValuesProducer valuesProducer)
             indexCreator = () -> createKNNIndexFromScratch(field, pair, knnEngine, indexPath);
         }
 
+        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
         // not be marked as added to the directory.
@@ -143,6 +172,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<String, Object> parameters = ImmutableMap.of(
             KNNConstants.INDEX_THREAD_QTY,
@@ -210,7 +257,12 @@ public void merge(MergeState mergeState) {
             for (FieldInfo fieldInfo : mergeState.mergeFieldInfos) {
                 DocValuesType type = fieldInfo.getDocValuesType();
                 if (type == DocValuesType.BINARY && fieldInfo.attributes().containsKey(KNNVectorFieldMapper.KNN_FIELD)) {
-                    addKNNBinaryField(fieldInfo, new KNN80DocValuesReader(mergeState));
+                    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) {

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);
+    }
+}

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<String, KNNStat<?>> buildStatsMap() {
         addEngineStats(builder);
         addScriptStats(builder);
         addModelStats(builder);
+        addGraphStats(builder);
         return builder.build();
     }
 
@@ -169,4 +171,31 @@ private void addModelStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
                 new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeAsPercentage))
             );
     }
+
+    private void addGraphStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
+        builder.put(StatNames.GRAPH_STATS.getName(), new KNNStat<>(false, new Supplier<Map<String, Map<String, Object>>>() {
+            @Override
+            public Map<String, Map<String, Object>> get() {
+                return createGraphStatsMap();
+            }
+        }));
+    }
+
+    private Map<String, Map<String, Object>> createGraphStatsMap() {
+        Map<String, Object> 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<String, Object> refreshMap = new HashMap<>();
+        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<String, Map<String, Object>> graphStatsMap = new HashMap<>();
+        graphStatsMap.put(StatNames.MERGE.getName(), mergeMap);
+        graphStatsMap.put(StatNames.REFRESH.getName(), refreshMap);
+        return graphStatsMap;
+    }
 }

src/main/java/org/opensearch/knn/plugin/stats/StatNames.java

@@ -41,7 +41,10 @@ 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"),
+    MERGE("merge");
 
     private String name;