update PR based on comments

Signed-off-by: Chenyang Ji <cyji@amazon.com>

Author: ansjcy

src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java

@@ -19,6 +19,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.stream.Collectors;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.opensearch.client.Client;
@@ -449,10 +450,9 @@ protected void doClose() throws IOException {
      * @return QueryInsightsHealthStats
      */
     public QueryInsightsHealthStats getHealthStats() {
-        Map<MetricType, TopQueriesHealthStats> topQueriesHealthStatsMap = new HashMap<>();
-        for (Map.Entry<MetricType, TopQueriesService> entry : topQueriesServices.entrySet()) {
-            topQueriesHealthStatsMap.put(entry.getKey(), entry.getValue().getHealthStats());
-        }
+        Map<MetricType, TopQueriesHealthStats> topQueriesHealthStatsMap = topQueriesServices.entrySet()
+            .stream()
+            .collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().getHealthStats()));
         return new QueryInsightsHealthStats(
             threadPool.info(QUERY_INSIGHTS_EXECUTOR),
             this.queryRecordsQueue.size(),

src/main/java/org/opensearch/plugin/insights/rules/model/healthStats/QueryGrouperHealthStats.java

@@ -21,8 +21,8 @@
 public class QueryGrouperHealthStats implements ToXContentFragment, Writeable {
     private final int queryGroupCount;
     private final int queryGroupHeapSize;
-    private static final String QUERY_GROUP_COUNT = "QueryGroupCount";
-    private static final String QUERY_GROUP_HEAP_SIZE = "QueryGroupHeapSize";
+    private static final String QUERY_GROUP_COUNT_TOTAL = "QueryGroupCount_Total";
+    private static final String QUERY_GROUP_COUNT_MAX_HEAP = "QueryGroupCount_MaxHeap";
 
     /**
      * Constructor to read QueryGrouperHealthStats from a StreamInput.
@@ -67,8 +67,8 @@ public void writeTo(StreamOutput out) throws IOException {
      */
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.field(QUERY_GROUP_COUNT, queryGroupCount);
-        builder.field(QUERY_GROUP_HEAP_SIZE, queryGroupHeapSize);
+        builder.field(QUERY_GROUP_COUNT_TOTAL, queryGroupCount);
+        builder.field(QUERY_GROUP_COUNT_MAX_HEAP, queryGroupHeapSize);
         return builder;
     }
 

src/test/java/org/opensearch/plugin/insights/core/service/TopQueriesServiceTests.java

@@ -149,7 +149,6 @@ public void testRollingWindowsWithDifferentGroup() {
     }
 
     public void testGetHealthStats_EmptyService() {
-        // Get the health stats from an empty TopQueriesService
         TopQueriesHealthStats healthStats = topQueriesService.getHealthStats();
         // Validate the health stats
         assertNotNull(healthStats);
@@ -160,16 +159,13 @@ public void testGetHealthStats_EmptyService() {
     }
 
     public void testGetHealthStats_WithData() {
-        // Add some mock records to the TopQueriesService
         List<SearchQueryRecord> records = QueryInsightsTestUtils.generateQueryInsightRecords(2);
         topQueriesService.consumeRecords(records);
-        // Get the health stats after adding data
         TopQueriesHealthStats healthStats = topQueriesService.getHealthStats();
-        // Validate the health stats
         assertNotNull(healthStats);
         assertEquals(2, healthStats.getTopQueriesHeapSize()); // Since we added two records
         assertNotNull(healthStats.getQueryGrouperHealthStats());
-        // Assuming no grouping by default, expect QueryGroupCount to be 2
+        // Assuming no grouping by default, expect QueryGroupCount to be 0
         assertEquals(0, healthStats.getQueryGrouperHealthStats().getQueryGroupCount());
     }
 }

src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java

@@ -694,7 +694,6 @@ public void testGetHealthStatsWithGroups() {
         QueryGrouperHealthStats healthStats = minMaxHeapQueryGrouper.getHealthStats();
         // Verify that group count stats reflect the correct number of total groups
         assertEquals(2, healthStats.getQueryGroupCount());
-        // Verify that heap size reflect the correct number of groups in max heap
         assertEquals(0, healthStats.getQueryGroupHeapSize());
     }
 

src/test/java/org/opensearch/plugin/insights/rules/model/healthStats/QueryGrouperHealthStatsTests.java

@@ -38,21 +38,19 @@ public void testSerialization() throws IOException {
         // Read from StreamInput
         StreamInput in = StreamInput.wrap(out.bytes().toBytesRef().bytes);
         QueryGrouperHealthStats deserializedStats = new QueryGrouperHealthStats(in);
-        // Assert equality
         assertEquals(stats.getQueryGroupCount(), deserializedStats.getQueryGroupCount());
         assertEquals(stats.getQueryGroupHeapSize(), deserializedStats.getQueryGroupHeapSize());
     }
 
     public void testToXContent() throws IOException {
         QueryGrouperHealthStats stats = new QueryGrouperHealthStats(queryGroupCount, queryGroupHeapSize);
-        // Write to XContent
         XContentBuilder builder = XContentFactory.jsonBuilder();
         builder.startObject();
         stats.toXContent(builder, ToXContent.EMPTY_PARAMS);
         builder.endObject();
         String expectedJson = String.format(
             Locale.ROOT,
-            "{\"QueryGroupCount\":%d,\"QueryGroupHeapSize\":%d}",
+            "{\"QueryGroupCount_Total\":%d,\"QueryGroupCount_MaxHeap\":%d}",
             queryGroupCount,
             queryGroupHeapSize
         );

src/test/java/org/opensearch/plugin/insights/rules/model/healthStats/QueryInsightsHealthStatsTests.java

@@ -69,15 +69,13 @@ public void testSerialization() throws IOException {
         // Read from StreamInput
         StreamInput in = StreamInput.wrap(out.bytes().toBytesRef().bytes);
         QueryInsightsHealthStats deserializedHealthStats = new QueryInsightsHealthStats(in);
-        // Assert equality
         assertEquals(healthStats.getQueryRecordsQueueSize(), deserializedHealthStats.getQueryRecordsQueueSize());
         assertNotNull(deserializedHealthStats.getThreadPoolInfo());
         assertNotNull(deserializedHealthStats.getTopQueriesHealthStats());
     }
 
     public void testToXContent() throws IOException {
         QueryInsightsHealthStats healthStats = new QueryInsightsHealthStats(threadPoolInfo, queryRecordsQueueSize, topQueriesHealthStats);
-        // Create XContentBuilder to build JSON
         XContentBuilder builder = XContentFactory.jsonBuilder();
         builder.startObject();
 
@@ -99,8 +97,8 @@ public void testToXContent() throws IOException {
             + "    \"TopQueriesHealthStats\": {\n"
             + "        \"latency\": {\n"
             + "            \"TopQueriesHeapSize\": 10,\n"
-            + "            \"QueryGroupCount\": 20,\n"
-            + "            \"QueryGroupHeapSize\": 15\n"
+            + "            \"QueryGroupCount_Total\": 20,\n"
+            + "            \"QueryGroupCount_MaxHeap\": 15\n"
             + "        }\n"
             + "    }\n"
             + "}";

src/test/java/org/opensearch/plugin/insights/rules/model/healthStats/TopQueriesHealthStatsTests.java

@@ -38,7 +38,6 @@ public void testSerialization() throws IOException {
         // Read from StreamInput
         StreamInput in = StreamInput.wrap(out.bytes().toBytesRef().bytes);
         TopQueriesHealthStats deserializedHealthStats = new TopQueriesHealthStats(in);
-        // Assert equality
         assertEquals(healthStats.getTopQueriesHeapSize(), deserializedHealthStats.getTopQueriesHeapSize());
         assertNotNull(deserializedHealthStats.getQueryGrouperHealthStats());
         assertEquals(
@@ -53,14 +52,13 @@ public void testSerialization() throws IOException {
 
     public void testToXContent() throws IOException {
         TopQueriesHealthStats healthStats = new TopQueriesHealthStats(topQueriesHeapSize, queryGrouperHealthStats);
-        // Write to XContent
         XContentBuilder builder = XContentFactory.jsonBuilder();
         builder.startObject();
         healthStats.toXContent(builder, ToXContent.EMPTY_PARAMS);
         builder.endObject();
         String expectedJson = String.format(
             Locale.ROOT,
-            "{\"TopQueriesHeapSize\":%d,\"QueryGroupCount\":%d,\"QueryGroupHeapSize\":%d}",
+            "{\"TopQueriesHeapSize\":%d,\"QueryGroupCount_Total\":%d,\"QueryGroupCount_MaxHeap\":%d}",
             topQueriesHeapSize,
             queryGrouperHealthStats.getQueryGroupCount(),
             queryGrouperHealthStats.getQueryGroupHeapSize()