From 2aa215be844305bdf9e501120abd2f5c2c237a61 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Mon, 18 Nov 2024 22:03:52 +0000 Subject: [PATCH] Add type attribute to search query record (#157) * Add type attribute to search query record Signed-off-by: Siddhant Deshmukh * Make type attribute enum Signed-off-by: Siddhant Deshmukh * Use group_by attribute Signed-off-by: Siddhant Deshmukh --------- Signed-off-by: Siddhant Deshmukh (cherry picked from commit b5e2d07fb2dcb8b1f9175ac826bcff042e9b02d9) Signed-off-by: github-actions[bot] --- .../core/listener/QueryInsightsListener.java | 5 ++--- .../service/grouper/MinMaxHeapQueryGrouper.java | 1 + .../plugin/insights/rules/model/Attribute.java | 10 +++++++++- .../insights/rules/model/SearchQueryRecord.java | 8 ++++++++ .../plugin/insights/QueryInsightsTestUtils.java | 2 ++ .../grouper/MinMaxHeapQueryGrouperTests.java | 14 ++++++++++++++ 6 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java index aa141f3f..34b9673d 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java +++ b/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 diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java b/src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java index d9c47741..19fb9a28 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java +++ b/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(); diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java b/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java index 8409796b..f5951e95 100644 --- a/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java +++ b/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) 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(); } diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java b/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java index 5877f676..66fe7675 100644 --- a/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java +++ b/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)); diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java index 980ee6a0..a7151933 100644 --- a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java +++ b/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 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( diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java index ec7ea309..ddea2ca1 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java +++ b/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 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)); + } + } }