Add type attribute to search query record

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

@@ -17,7 +17,6 @@
 import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getTopNSizeSetting;
 import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getTopNWindowSizeSetting;
 
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
@@ -33,7 +32,6 @@
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
 import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo;
-import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.plugin.insights.core.metrics.OperationalMetric;
 import org.opensearch.plugin.insights.core.metrics.OperationalMetricsCounter;
 import org.opensearch.plugin.insights.core.service.QueryInsightsService;
@@ -42,6 +40,7 @@
 import org.opensearch.plugin.insights.rules.model.Measurement;
 import org.opensearch.plugin.insights.rules.model.MetricType;
 import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
+import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
 import org.opensearch.tasks.Task;
 
 /**
@@ -50,7 +49,6 @@
  * either for each request or for each phase.
  */
 public final class QueryInsightsListener extends SearchRequestOperationsListener {
-    private static final ToXContent.Params FORMAT_PARAMS = new ToXContent.MapParams(Collections.singletonMap("pretty", "false"));
 
     private static final Logger log = LogManager.getLogger(QueryInsightsListener.class);
 
@@ -271,6 +269,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.GROUP_BY, QueryInsightsSettings.DEFAULT_GROUPING_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

@@ -127,6 +127,7 @@ public SearchQueryRecord add(final SearchQueryRecord searchQueryRecord) {
             aggregateSearchQueryRecord = searchQueryRecord;
             aggregateSearchQueryRecord.setGroupingId(groupId);
             aggregateSearchQueryRecord.setMeasurementAggregation(metricType, aggregationType);
+            aggregateSearchQueryRecord.addAttribute(Attribute.GROUP_BY, groupingType);
             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,
+    /**
+     * Grouping type of the query record (none, similarity)
+     */
+    GROUP_BY;
 
     /**
      * Read an Attribute from a StreamInput
@@ -97,6 +101,8 @@ public static void writeValueTo(StreamOutput out, Object attributeValue) throws
             out.writeList((List<? extends Writeable>) attributeValue);
         } else if (attributeValue instanceof SearchSourceBuilder) {
             ((SearchSourceBuilder) attributeValue).writeTo(out);
+        } else if (attributeValue instanceof GroupingType) {
+            out.writeString(((GroupingType) attributeValue).getValue());
         } else {
             out.writeGenericValue(attributeValue);
         }
@@ -116,6 +122,8 @@ public static Object readAttributeValue(StreamInput in, Attribute attribute) thr
         } else if (attribute == Attribute.SOURCE) {
             SearchSourceBuilder builder = new SearchSourceBuilder(in);
             return builder;
+        } else if (attribute == Attribute.GROUP_BY) {
+            return GroupingType.valueOf(in.readString().toUpperCase(Locale.ROOT));
         } else {
             return in.readGenericValue();
         }

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

@@ -8,6 +8,8 @@
 
 package org.opensearch.plugin.insights.rules.model;
 
+import static org.opensearch.plugin.insights.rules.model.Attribute.GROUP_BY;
+
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -89,6 +91,10 @@ public class SearchQueryRecord implements ToXContentObject, Writeable {
      * Custom search request labels
      */
     public static final String LABELS = "labels";
+    /**
+     * Grouping type of the query record (none, similarity)
+     */
+    public static final String GROUP_BY = "group_by";
 
     public static final String MEASUREMENTS = "measurements";
     private String groupingId;
@@ -167,6 +173,8 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc
                     case SEARCH_TYPE:
                         attributes.put(Attribute.SEARCH_TYPE, parser.text());
                         break;
+                    case GROUP_BY:
+                        attributes.put(Attribute.GROUP_BY, 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

@@ -43,6 +43,7 @@
 import org.opensearch.plugin.insights.rules.action.top_queries.TopQueries;
 import org.opensearch.plugin.insights.rules.model.AggregationType;
 import org.opensearch.plugin.insights.rules.model.Attribute;
+import org.opensearch.plugin.insights.rules.model.GroupingType;
 import org.opensearch.plugin.insights.rules.model.Measurement;
 import org.opensearch.plugin.insights.rules.model.MetricType;
 import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
@@ -139,6 +140,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.GROUP_BY, GroupingType.NONE);
             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(GroupingType.NONE, record.getAttributes().get(Attribute.GROUP_BY));
+        }
+        SearchQueryRecord groupedRecord;
+        for (SearchQueryRecord record : records) {
+            groupedRecord = minMaxHeapQueryGrouper.add(record);
+            assertEquals(GroupingType.SIMILARITY, groupedRecord.getAttributes().get(Attribute.GROUP_BY));
+        }
+    }
 }