Use group_by attribute

Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>

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,7 +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.rules.model.Type;
+import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
 import org.opensearch.tasks.Task;
 
 /**
@@ -51,8 +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 Type DEFAULT_TYPE = Type.query;
 
     private static final Logger log = LogManager.getLogger(QueryInsightsListener.class);
 
@@ -273,7 +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.TYPE, DEFAULT_TYPE);
+            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

@@ -20,7 +20,6 @@
 import org.opensearch.plugin.insights.rules.model.GroupingType;
 import org.opensearch.plugin.insights.rules.model.MetricType;
 import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
-import org.opensearch.plugin.insights.rules.model.Type;
 import org.opensearch.plugin.insights.rules.model.healthStats.QueryGrouperHealthStats;
 import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
 
@@ -128,7 +127,7 @@ public SearchQueryRecord add(final SearchQueryRecord searchQueryRecord) {
             aggregateSearchQueryRecord = searchQueryRecord;
             aggregateSearchQueryRecord.setGroupingId(groupId);
             aggregateSearchQueryRecord.setMeasurementAggregation(metricType, aggregationType);
-            aggregateSearchQueryRecord.addAttribute(Attribute.TYPE, Type.group);
+            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

@@ -62,9 +62,9 @@ public enum Attribute {
      */
     QUERY_HASHCODE,
     /**
-     * Type of the query record (either 'query' or 'group')
+     * Grouping type of the query record (none, similarity)
      */
-    TYPE;
+    GROUP_BY;
 
     /**
      * Read an Attribute from a StreamInput
@@ -101,10 +101,10 @@ 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 Type) {
-            out.writeString(((Type) attributeValue).getValue());
+        } else if (attributeValue instanceof GroupingType) {
+            out.writeString(((GroupingType) attributeValue).getValue());
         } else {
-        out.writeGenericValue(attributeValue);
+            out.writeGenericValue(attributeValue);
         }
     }
 
@@ -122,8 +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.TYPE) {
-            return Type.valueOf(in.readString());
+        } 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;
@@ -90,9 +92,9 @@ public class SearchQueryRecord implements ToXContentObject, Writeable {
      */
     public static final String LABELS = "labels";
     /**
-     * Type of the query record (either 'query' or 'group')
+     * Grouping type of the query record (none, similarity)
      */
-    public static final String TYPE = "type";
+    public static final String GROUP_BY = "group_by";
 
     public static final String MEASUREMENTS = "measurements";
     private String groupingId;
@@ -171,8 +173,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 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,10 +43,10 @@
 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;
-import org.opensearch.plugin.insights.rules.model.Type;
 import org.opensearch.plugin.insights.settings.QueryCategorizationSettings;
 import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
 import org.opensearch.search.builder.SearchSourceBuilder;
@@ -140,7 +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.TYPE, Type.query);
+            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

@@ -19,7 +19,6 @@
 import org.opensearch.plugin.insights.rules.model.GroupingType;
 import org.opensearch.plugin.insights.rules.model.MetricType;
 import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
-import org.opensearch.plugin.insights.rules.model.Type;
 import org.opensearch.plugin.insights.rules.model.healthStats.QueryGrouperHealthStats;
 import org.opensearch.test.OpenSearchTestCase;
 
@@ -707,12 +706,12 @@ public void testAttributeTypeSetToGroup() {
         final List<SearchQueryRecord> records = QueryInsightsTestUtils.generateQueryInsightRecords(numOfRecords);
 
         for (SearchQueryRecord record : records) {
-            assertEquals(Type.query, record.getAttributes().get(Attribute.TYPE));
+            assertEquals(GroupingType.NONE, record.getAttributes().get(Attribute.GROUP_BY));
         }
         SearchQueryRecord groupedRecord;
         for (SearchQueryRecord record : records) {
             groupedRecord = minMaxHeapQueryGrouper.add(record);
-            assertEquals(Type.group, groupedRecord.getAttributes().get(Attribute.TYPE));
+            assertEquals(GroupingType.SIMILARITY, groupedRecord.getAttributes().get(Attribute.GROUP_BY));
         }
     }
 }