Refactor stat files for better readability.

CHANGELOG.md

@@ -28,3 +28,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Fixing Lucene912Codec Issue with BWC for Lucene 10.0.1 upgrade[#2429](https://github.com/opensearch-project/k-NN/pull/2429)
 * Enabled idempotency of local builds when using `./gradlew clean` and nest `jni/release` directory under `jni/build` for easier cleanup [#2516](https://github.com/opensearch-project/k-NN/pull/2516)
 ### Refactoring
+* Refactored the KNN Stat files for better readability.

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

@@ -55,7 +55,7 @@ public Map<String, KNNStat<?>> getStats() {
      * @return Map of stats kept at the node level
      */
     public Map<String, KNNStat<?>> getNodeStats() {
-        return getClusterOrNodeStats(false);
+        return getFilteredStats(false);
     }
 
     /**
@@ -64,18 +64,14 @@ public Map<String, KNNStat<?>> getNodeStats() {
      * @return Map of stats kept at the cluster level
      */
     public Map<String, KNNStat<?>> getClusterStats() {
-        return getClusterOrNodeStats(true);
+        return getFilteredStats(true);
     }
 
-    private Map<String, KNNStat<?>> getClusterOrNodeStats(Boolean getClusterStats) {
-        Map<String, KNNStat<?>> statsMap = new HashMap<>();
-
-        for (Map.Entry<String, KNNStat<?>> entry : knnStats.entrySet()) {
-            if (entry.getValue().isClusterLevel() == getClusterStats) {
-                statsMap.put(entry.getKey(), entry.getValue());
-            }
-        }
-        return statsMap;
+    private Map<String, KNNStat<?>> getFilteredStats(boolean isClusterLevel) {
+        return knnStats.entrySet()
+            .stream()
+            .filter(entry -> entry.getValue().isClusterLevel() == isClusterLevel)
+            .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue));
     }
 
     private Map<String, KNNStat<?>> buildStatsMap() {
@@ -91,109 +87,105 @@ private Map<String, KNNStat<?>> buildStatsMap() {
 
     private void addQueryStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
         // KNN Query Stats
-        builder.put(StatNames.KNN_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_REQUESTS)))
+        builder.put(StatNames.KNN_QUERY_REQUESTS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.KNN_QUERY_REQUESTS)))
             .put(
                 StatNames.KNN_QUERY_WITH_FILTER_REQUESTS.getName(),
-                new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS))
+                createNodeStat(new KNNCounterSupplier(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS))
             );
 
         // Min Score Query Stats
         builder.put(
             StatNames.MIN_SCORE_QUERY_REQUESTS.getName(),
-            new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MIN_SCORE_QUERY_REQUESTS))
+            createNodeStat(new KNNCounterSupplier(KNNCounter.MIN_SCORE_QUERY_REQUESTS))
         )
             .put(
                 StatNames.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS.getName(),
-                new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS))
+                createNodeStat(new KNNCounterSupplier(KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS))
             );
 
         // Max Distance Query Stats
         builder.put(
             StatNames.MAX_DISTANCE_QUERY_REQUESTS.getName(),
-            new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MAX_DISTANCE_QUERY_REQUESTS))
+            createNodeStat(new KNNCounterSupplier(KNNCounter.MAX_DISTANCE_QUERY_REQUESTS))
         )
             .put(
                 StatNames.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS.getName(),
-                new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS))
+                createNodeStat(new KNNCounterSupplier(KNNCounter.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS))
             );
     }
 
     private void addNativeMemoryStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
-        builder.put(StatNames.HIT_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::hitCount)))
-            .put(StatNames.MISS_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::missCount)))
-            .put(StatNames.LOAD_SUCCESS_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::loadSuccessCount)))
-            .put(
-                StatNames.LOAD_EXCEPTION_COUNT.getName(),
-                new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::loadExceptionCount))
-            )
-            .put(StatNames.TOTAL_LOAD_TIME.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::totalLoadTime)))
-            .put(StatNames.EVICTION_COUNT.getName(), new KNNStat<>(false, new KNNInnerCacheStatsSupplier(CacheStats::evictionCount)))
+        builder.put(StatNames.HIT_COUNT.getName(), createNodeStat(new KNNInnerCacheStatsSupplier(CacheStats::hitCount)))
+            .put(StatNames.MISS_COUNT.getName(), createNodeStat(new KNNInnerCacheStatsSupplier(CacheStats::missCount)))
+            .put(StatNames.LOAD_SUCCESS_COUNT.getName(), createNodeStat(new KNNInnerCacheStatsSupplier(CacheStats::loadSuccessCount)))
+            .put(StatNames.LOAD_EXCEPTION_COUNT.getName(), createNodeStat(new KNNInnerCacheStatsSupplier(CacheStats::loadExceptionCount)))
+            .put(StatNames.TOTAL_LOAD_TIME.getName(), createNodeStat(new KNNInnerCacheStatsSupplier(CacheStats::totalLoadTime)))
+            .put(StatNames.EVICTION_COUNT.getName(), createNodeStat(new KNNInnerCacheStatsSupplier(CacheStats::evictionCount)))
             .put(
                 StatNames.GRAPH_MEMORY_USAGE.getName(),
-                new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeInKilobytes))
+                createNodeStat(new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeInKilobytes))
             )
             .put(
                 StatNames.GRAPH_MEMORY_USAGE_PERCENTAGE.getName(),
-                new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeAsPercentage))
+                createNodeStat(new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesSizeAsPercentage))
             )
             .put(
                 StatNames.INDICES_IN_CACHE.getName(),
-                new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesCacheStats))
+                createNodeStat(new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getIndicesCacheStats))
             )
             .put(
                 StatNames.CACHE_CAPACITY_REACHED.getName(),
-                new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::isCacheCapacityReached))
+                createNodeStat(new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::isCacheCapacityReached))
             )
-            .put(StatNames.GRAPH_QUERY_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_ERRORS)))
-            .put(StatNames.GRAPH_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_REQUESTS)))
-            .put(StatNames.GRAPH_INDEX_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_ERRORS)))
-            .put(StatNames.GRAPH_INDEX_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_REQUESTS)))
-            .put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), new KNNStat<>(true, new KNNCircuitBreakerSupplier()));
+            .put(StatNames.GRAPH_QUERY_ERRORS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_ERRORS)))
+            .put(StatNames.GRAPH_QUERY_REQUESTS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.GRAPH_QUERY_REQUESTS)))
+            .put(StatNames.GRAPH_INDEX_ERRORS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_ERRORS)))
+            .put(StatNames.GRAPH_INDEX_REQUESTS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.GRAPH_INDEX_REQUESTS)))
+            .put(StatNames.CIRCUIT_BREAKER_TRIGGERED.getName(), createClusterStat(new KNNCircuitBreakerSupplier()));
     }
 
     private void addEngineStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
-        builder.put(StatNames.FAISS_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.FAISS)))
-            .put(StatNames.NMSLIB_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.NMSLIB)))
-            .put(StatNames.LUCENE_LOADED.getName(), new KNNStat<>(false, new LibraryInitializedSupplier(KNNEngine.LUCENE)));
+        builder.put(StatNames.FAISS_LOADED.getName(), createNodeStat(new LibraryInitializedSupplier(KNNEngine.FAISS)))
+            .put(StatNames.NMSLIB_LOADED.getName(), createNodeStat(new LibraryInitializedSupplier(KNNEngine.NMSLIB)))
+            .put(StatNames.LUCENE_LOADED.getName(), createNodeStat(new LibraryInitializedSupplier(KNNEngine.LUCENE)));
     }
 
     private void addScriptStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
-        builder.put(StatNames.SCRIPT_COMPILATIONS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATIONS)))
+        builder.put(StatNames.SCRIPT_COMPILATIONS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATIONS)))
             .put(
                 StatNames.SCRIPT_COMPILATION_ERRORS.getName(),
-                new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATION_ERRORS))
+                createNodeStat(new KNNCounterSupplier(KNNCounter.SCRIPT_COMPILATION_ERRORS))
             )
-            .put(StatNames.SCRIPT_QUERY_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_QUERY_REQUESTS)))
-            .put(StatNames.SCRIPT_QUERY_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.SCRIPT_QUERY_ERRORS)));
+            .put(StatNames.SCRIPT_QUERY_REQUESTS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.SCRIPT_QUERY_REQUESTS)))
+            .put(StatNames.SCRIPT_QUERY_ERRORS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.SCRIPT_QUERY_ERRORS)));
     }
 
     private void addModelStats(ImmutableMap.Builder<String, KNNStat<?>> builder) {
         builder.put(
             StatNames.INDEXING_FROM_MODEL_DEGRADED.getName(),
-            new KNNStat<>(
-                false,
+            createNodeStat(
                 new EventOccurredWithinThresholdSupplier(
                     new ModelIndexingDegradingSupplier(ModelCache::getEvictedDueToSizeAt),
                     KNNConstants.MODEL_CACHE_CAPACITY_ATROPHY_THRESHOLD_IN_MINUTES,
                     ChronoUnit.MINUTES
                 )
             )
         )
-            .put(StatNames.MODEL_INDEX_STATUS.getName(), new KNNStat<>(true, new ModelIndexStatusSupplier<>(ModelDao::getHealthStatus)))
-            .put(StatNames.TRAINING_REQUESTS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.TRAINING_REQUESTS)))
-            .put(StatNames.TRAINING_ERRORS.getName(), new KNNStat<>(false, new KNNCounterSupplier(KNNCounter.TRAINING_ERRORS)))
+            .put(StatNames.MODEL_INDEX_STATUS.getName(), createClusterStat(new ModelIndexStatusSupplier<>(ModelDao::getHealthStatus)))
+            .put(StatNames.TRAINING_REQUESTS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.TRAINING_REQUESTS)))
+            .put(StatNames.TRAINING_ERRORS.getName(), createNodeStat(new KNNCounterSupplier(KNNCounter.TRAINING_ERRORS)))
             .put(
                 StatNames.TRAINING_MEMORY_USAGE.getName(),
-                new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeInKilobytes))
+                createNodeStat(new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeInKilobytes))
             )
             .put(
                 StatNames.TRAINING_MEMORY_USAGE_PERCENTAGE.getName(),
-                new KNNStat<>(false, new NativeMemoryCacheManagerSupplier<>(NativeMemoryCacheManager::getTrainingSizeAsPercentage))
+                createNodeStat(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>>>() {
+        builder.put(StatNames.GRAPH_STATS.getName(), createNodeStat(new Supplier<Map<String, Map<String, Object>>>() {
             @Override
             public Map<String, Map<String, Object>> get() {
                 return createGraphStatsMap();
@@ -218,4 +210,12 @@ private Map<String, Map<String, Object>> createGraphStatsMap() {
         graphStatsMap.put(StatNames.REFRESH.getName(), refreshMap);
         return graphStatsMap;
     }
+
+    private static <T> KNNStat<T> createNodeStat(Supplier<T> supplier) {
+        return new KNNStat<>(false, supplier);
+    }
+
+    private static <T> KNNStat<T> createClusterStat(Supplier<T> supplier) {
+        return new KNNStat<>(true, supplier);
+    }
 }

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

@@ -12,41 +12,55 @@
  * Enum contains names of the stats
  */
 public enum StatNames {
+    // Cache stats
     HIT_COUNT("hit_count"),
     MISS_COUNT("miss_count"),
     LOAD_SUCCESS_COUNT("load_success_count"),
     LOAD_EXCEPTION_COUNT("load_exception_count"),
     TOTAL_LOAD_TIME("total_load_time"),
     EVICTION_COUNT("eviction_count"),
+
     GRAPH_MEMORY_USAGE("graph_memory_usage"),
     GRAPH_MEMORY_USAGE_PERCENTAGE("graph_memory_usage_percentage"),
     CACHE_CAPACITY_REACHED("cache_capacity_reached"),
     INDICES_IN_CACHE("indices_in_cache"),
+
     CIRCUIT_BREAKER_TRIGGERED("circuit_breaker_triggered"),
-    MODEL_INDEX_STATUS("model_index_status"),
+
+    // Engine Stats
     FAISS_LOADED("faiss_initialized"),
     NMSLIB_LOADED("nmslib_initialized"),
     LUCENE_LOADED("lucene_initialized"),
-    INDEXING_FROM_MODEL_DEGRADED("indexing_from_model_degraded"),
+
     GRAPH_QUERY_ERRORS(KNNCounter.GRAPH_QUERY_ERRORS.getName()),
     GRAPH_QUERY_REQUESTS(KNNCounter.GRAPH_QUERY_REQUESTS.getName()),
     GRAPH_INDEX_ERRORS(KNNCounter.GRAPH_INDEX_ERRORS.getName()),
     GRAPH_INDEX_REQUESTS(KNNCounter.GRAPH_INDEX_REQUESTS.getName()),
-    KNN_QUERY_REQUESTS(KNNCounter.KNN_QUERY_REQUESTS.getName()),
+
+    // Script Stats
     SCRIPT_COMPILATIONS(KNNCounter.SCRIPT_COMPILATIONS.getName()),
     SCRIPT_COMPILATION_ERRORS(KNNCounter.SCRIPT_COMPILATION_ERRORS.getName()),
     SCRIPT_QUERY_REQUESTS(KNNCounter.SCRIPT_QUERY_REQUESTS.getName()),
+    SCRIPT_QUERY_ERRORS(KNNCounter.SCRIPT_QUERY_ERRORS.getName()),
+
+    // Model Stats
+    MODEL_INDEX_STATUS("model_index_status"),
+    INDEXING_FROM_MODEL_DEGRADED("indexing_from_model_degraded"),
     TRAINING_REQUESTS(KNNCounter.TRAINING_REQUESTS.getName()),
     TRAINING_ERRORS(KNNCounter.TRAINING_ERRORS.getName()),
     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()),
     GRAPH_STATS("graph_stats"),
     REFRESH("refresh"),
     MERGE("merge"),
+
+    // KNN Query Stats
+    KNN_QUERY_REQUESTS(KNNCounter.KNN_QUERY_REQUESTS.getName()),
+    KNN_QUERY_WITH_FILTER_REQUESTS(KNNCounter.KNN_QUERY_WITH_FILTER_REQUESTS.getName()),
+
     MIN_SCORE_QUERY_REQUESTS(KNNCounter.MIN_SCORE_QUERY_REQUESTS.getName()),
     MIN_SCORE_QUERY_WITH_FILTER_REQUESTS(KNNCounter.MIN_SCORE_QUERY_WITH_FILTER_REQUESTS.getName()),
+
     MAX_DISTANCE_QUERY_REQUESTS(KNNCounter.MAX_DISTANCE_QUERY_REQUESTS.getName()),
     MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS(KNNCounter.MAX_DISTANCE_QUERY_WITH_FILTER_REQUESTS.getName());