Add type attribute to search query record

src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java

@@ -51,6 +51,7 @@
  */
 public final class QueryInsightsListener extends SearchRequestOperationsListener {
     private static final ToXContent.Params FORMAT_PARAMS = new ToXContent.MapParams(Collections.singletonMap("pretty", "false"));
+    private static final String DEFAULT_TYPE = "query";
 
     private static final Logger log = LogManager.getLogger(QueryInsightsListener.class);
 
@@ -271,6 +272,7 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final
             attributes.put(Attribute.INDICES, request.indices());
             attributes.put(Attribute.PHASE_LATENCY_MAP, searchRequestContext.phaseTookMap());
             attributes.put(Attribute.TASK_RESOURCE_USAGES, tasksResourceUsages);
+            attributes.put(Attribute.TYPE, DEFAULT_TYPE);
 
             if (queryInsightsService.isGroupingEnabled() || log.isTraceEnabled()) {
                 // Generate the query shape only if grouping is enabled or trace logging is enabled

src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java

@@ -29,6 +29,8 @@
  */
 public class MinMaxHeapQueryGrouper implements QueryGrouper {
 
+    private static final String GROUP_TYPE = "group";
+
     /**
      * Logger
      */
@@ -127,6 +129,7 @@ public SearchQueryRecord add(final SearchQueryRecord searchQueryRecord) {
             aggregateSearchQueryRecord = searchQueryRecord;
             aggregateSearchQueryRecord.setGroupingId(groupId);
             aggregateSearchQueryRecord.setMeasurementAggregation(metricType, aggregationType);
+            aggregateSearchQueryRecord.addAttribute(Attribute.TYPE, GROUP_TYPE);
             addToMinPQ(aggregateSearchQueryRecord, groupId);
         } else {
             aggregateSearchQueryRecord = groupIdToAggSearchQueryRecord.get(groupId).v1();

src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java

@@ -60,7 +60,11 @@ public enum Attribute {
     /**
      * Unique hashcode used to group similar queries
      */
-    QUERY_HASHCODE;
+    QUERY_HASHCODE,
+    /**
+     * Type of the query record (either 'query' or 'group')
+     */
+    TYPE;
 
     /**
      * Read an Attribute from a StreamInput

src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java

@@ -89,6 +89,10 @@ public class SearchQueryRecord implements ToXContentObject, Writeable {
      * Custom search request labels
      */
     public static final String LABELS = "labels";
+    /**
+     * Type of the query record (either 'query' or 'group')
+     */
+    public static final String TYPE = "type";
 
     public static final String MEASUREMENTS = "measurements";
     private String groupingId;
@@ -167,6 +171,8 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc
                     case SEARCH_TYPE:
                         attributes.put(Attribute.SEARCH_TYPE, parser.text());
                         break;
+                    case TYPE:
+                        attributes.put(Attribute.TYPE, parser.text());
                     case SOURCE:
                         XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
                         attributes.put(Attribute.SOURCE, SearchSourceBuilder.fromXContent(parser, false));

src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java

@@ -139,6 +139,7 @@ public static List<SearchQueryRecord> generateQueryInsightRecords(
             attributes.put(Attribute.INDICES, randomArray(1, 3, Object[]::new, () -> randomAlphaOfLengthBetween(5, 10)));
             attributes.put(Attribute.PHASE_LATENCY_MAP, phaseLatencyMap);
             attributes.put(Attribute.QUERY_HASHCODE, Objects.hashCode(i));
+            attributes.put(Attribute.TYPE, "query");
             attributes.put(
                 Attribute.TASK_RESOURCE_USAGES,
                 List.of(

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

@@ -700,4 +700,18 @@ public void testGetHealthStatsWithGroups() {
     private MinMaxHeapQueryGrouper getQueryGroupingService(AggregationType aggregationType, int topNSize) {
         return new MinMaxHeapQueryGrouper(MetricType.LATENCY, GroupingType.SIMILARITY, aggregationType, topQueriesStore, topNSize);
     }
+
+    public void testAttributeTypeSetToGroup() {
+        int numOfRecords = 10;
+        final List<SearchQueryRecord> records = QueryInsightsTestUtils.generateQueryInsightRecords(numOfRecords);
+
+        for (SearchQueryRecord record : records) {
+            assertEquals("query", record.getAttributes().get(Attribute.TYPE));
+        }
+        SearchQueryRecord groupedRecord;
+        for (SearchQueryRecord record : records) {
+            groupedRecord = minMaxHeapQueryGrouper.add(record);
+            assertEquals("group", groupedRecord.getAttributes().get(Attribute.TYPE));
+        }
+    }
 }