From 3f6a00338d799429ceb1c7a7e1bfd3180dae56e7 Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Wed, 3 Jan 2024 12:28:23 -0800 Subject: [PATCH] Address review comments, refactor code Signed-off-by: Siddhant Deshmukh --- .../SearchQueryAggregationCategorizer.java | 82 ++++++++-------- .../SearchQueryCategorizingVisitor.java | 74 +-------------- .../action/search/SearchQueryCounters.java | 94 +++++++++++++++++-- .../search/SearchQueryCategorizerTests.java | 6 +- 4 files changed, 126 insertions(+), 130 deletions(-) diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java b/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java index 44585dc0cfd77..35eaa027903c5 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java +++ b/server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java @@ -46,7 +46,6 @@ import java.util.Collection; import java.util.HashMap; import java.util.Map; -import java.util.function.Function; /** * Increments the counters related to Aggregation Search Queries. @@ -55,50 +54,49 @@ public class SearchQueryAggregationCategorizer { private static final String TYPE_TAG = "type"; private final SearchQueryCounters searchQueryCounters; - private final Map, Function> aggregationHandlers; + private final Map, String> aggregationHandlers; public SearchQueryAggregationCategorizer(SearchQueryCounters searchQueryCounters) { this.searchQueryCounters = searchQueryCounters; - this.aggregationHandlers = initializeAggregationHandlers(); - } + this.aggregationHandlers = new HashMap<>(); - private Map, Function> initializeAggregationHandlers() { - Map, Function> handlers = new HashMap<>(); + initializeAggregationHandlers(); + } - handlers.put(TermsAggregationBuilder.class, aggregationBuilder -> TermsAggregationBuilder.NAME); - handlers.put(AvgAggregationBuilder.class, aggregationBuilder -> AvgAggregationBuilder.NAME); - handlers.put(SumAggregationBuilder.class, aggregationBuilder -> SumAggregationBuilder.NAME); - handlers.put(MaxAggregationBuilder.class, aggregationBuilder -> MaxAggregationBuilder.NAME); - handlers.put(MinAggregationBuilder.class, aggregationBuilder -> MinAggregationBuilder.NAME); - handlers.put(ScriptedMetricAggregationBuilder.class, aggregationBuilder -> ScriptedMetricAggregationBuilder.NAME); - handlers.put(ExtendedStatsAggregationBuilder.class, aggregationBuilder -> ExtendedStatsAggregationBuilder.NAME); - handlers.put(FilterAggregationBuilder.class, aggregationBuilder -> FilterAggregationBuilder.NAME); - handlers.put(FiltersAggregationBuilder.class, aggregationBuilder -> FiltersAggregationBuilder.NAME); - handlers.put(AdjacencyMatrixAggregationBuilder.class, aggregationBuilder -> AdjacencyMatrixAggregationBuilder.NAME); - handlers.put(SamplerAggregationBuilder.class, aggregationBuilder -> SamplerAggregationBuilder.NAME); - handlers.put(DiversifiedAggregationBuilder.class, aggregationBuilder -> DiversifiedAggregationBuilder.NAME); - handlers.put(GlobalAggregationBuilder.class, aggregationBuilder -> GlobalAggregationBuilder.NAME); - handlers.put(MissingAggregationBuilder.class, aggregationBuilder -> MissingAggregationBuilder.NAME); - handlers.put(NestedAggregationBuilder.class, aggregationBuilder -> NestedAggregationBuilder.NAME); - handlers.put(ReverseNestedAggregationBuilder.class, aggregationBuilder -> ReverseNestedAggregationBuilder.NAME); - handlers.put(GeoDistanceAggregationBuilder.class, aggregationBuilder -> GeoDistanceAggregationBuilder.NAME); - handlers.put(HistogramAggregationBuilder.class, aggregationBuilder -> HistogramAggregationBuilder.NAME); - handlers.put(SignificantTermsAggregationBuilder.class, aggregationBuilder -> SignificantTermsAggregationBuilder.NAME); - handlers.put(SignificantTextAggregationBuilder.class, aggregationBuilder -> SignificantTextAggregationBuilder.NAME); - handlers.put(DateHistogramAggregationBuilder.class, aggregationBuilder -> DateHistogramAggregationBuilder.NAME); - handlers.put(RangeAggregationBuilder.class, aggregationBuilder -> RangeAggregationBuilder.NAME); - handlers.put(DateRangeAggregationBuilder.class, aggregationBuilder -> DateRangeAggregationBuilder.NAME); - handlers.put(IpRangeAggregationBuilder.class, aggregationBuilder -> IpRangeAggregationBuilder.NAME); - handlers.put(PercentilesAggregationBuilder.class, aggregationBuilder -> PercentilesAggregationBuilder.NAME); - handlers.put(PercentileRanksAggregationBuilder.class, aggregationBuilder -> PercentileRanksAggregationBuilder.NAME); - handlers.put(MedianAbsoluteDeviationAggregationBuilder.class, aggregationBuilder -> MedianAbsoluteDeviationAggregationBuilder.NAME); - handlers.put(CardinalityAggregationBuilder.class, aggregationBuilder -> CardinalityAggregationBuilder.NAME); - handlers.put(TopHitsAggregationBuilder.class, aggregationBuilder -> TopHitsAggregationBuilder.NAME); - handlers.put(GeoCentroidAggregationBuilder.class, aggregationBuilder -> GeoCentroidAggregationBuilder.NAME); - handlers.put(CompositeAggregationBuilder.class, aggregationBuilder -> CompositeAggregationBuilder.NAME); - handlers.put(MultiTermsAggregationBuilder.class, aggregationBuilder -> MultiTermsAggregationBuilder.NAME); + private void initializeAggregationHandlers() { - return handlers; + aggregationHandlers.put(TermsAggregationBuilder.class, TermsAggregationBuilder.NAME); + aggregationHandlers.put(AvgAggregationBuilder.class, AvgAggregationBuilder.NAME); + aggregationHandlers.put(SumAggregationBuilder.class, SumAggregationBuilder.NAME); + aggregationHandlers.put(MaxAggregationBuilder.class, MaxAggregationBuilder.NAME); + aggregationHandlers.put(MinAggregationBuilder.class, MinAggregationBuilder.NAME); + aggregationHandlers.put(ScriptedMetricAggregationBuilder.class, ScriptedMetricAggregationBuilder.NAME); + aggregationHandlers.put(ExtendedStatsAggregationBuilder.class, ExtendedStatsAggregationBuilder.NAME); + aggregationHandlers.put(FilterAggregationBuilder.class, FilterAggregationBuilder.NAME); + aggregationHandlers.put(FiltersAggregationBuilder.class, FiltersAggregationBuilder.NAME); + aggregationHandlers.put(AdjacencyMatrixAggregationBuilder.class, AdjacencyMatrixAggregationBuilder.NAME); + aggregationHandlers.put(SamplerAggregationBuilder.class, SamplerAggregationBuilder.NAME); + aggregationHandlers.put(DiversifiedAggregationBuilder.class, DiversifiedAggregationBuilder.NAME); + aggregationHandlers.put(GlobalAggregationBuilder.class, GlobalAggregationBuilder.NAME); + aggregationHandlers.put(MissingAggregationBuilder.class, MissingAggregationBuilder.NAME); + aggregationHandlers.put(NestedAggregationBuilder.class, NestedAggregationBuilder.NAME); + aggregationHandlers.put(ReverseNestedAggregationBuilder.class, ReverseNestedAggregationBuilder.NAME); + aggregationHandlers.put(GeoDistanceAggregationBuilder.class, GeoDistanceAggregationBuilder.NAME); + aggregationHandlers.put(HistogramAggregationBuilder.class, HistogramAggregationBuilder.NAME); + aggregationHandlers.put(SignificantTermsAggregationBuilder.class, SignificantTermsAggregationBuilder.NAME); + aggregationHandlers.put(SignificantTextAggregationBuilder.class, SignificantTextAggregationBuilder.NAME); + aggregationHandlers.put(DateHistogramAggregationBuilder.class, DateHistogramAggregationBuilder.NAME); + aggregationHandlers.put(RangeAggregationBuilder.class, RangeAggregationBuilder.NAME); + aggregationHandlers.put(DateRangeAggregationBuilder.class, DateRangeAggregationBuilder.NAME); + aggregationHandlers.put(IpRangeAggregationBuilder.class, IpRangeAggregationBuilder.NAME); + aggregationHandlers.put(PercentilesAggregationBuilder.class, PercentilesAggregationBuilder.NAME); + aggregationHandlers.put(PercentileRanksAggregationBuilder.class, PercentileRanksAggregationBuilder.NAME); + aggregationHandlers.put(MedianAbsoluteDeviationAggregationBuilder.class, MedianAbsoluteDeviationAggregationBuilder.NAME); + aggregationHandlers.put(CardinalityAggregationBuilder.class, CardinalityAggregationBuilder.NAME); + aggregationHandlers.put(TopHitsAggregationBuilder.class, TopHitsAggregationBuilder.NAME); + aggregationHandlers.put(GeoCentroidAggregationBuilder.class, GeoCentroidAggregationBuilder.NAME); + aggregationHandlers.put(CompositeAggregationBuilder.class, CompositeAggregationBuilder.NAME); + aggregationHandlers.put(MultiTermsAggregationBuilder.class, MultiTermsAggregationBuilder.NAME); } public void incrementSearchQueryAggregationCounters(Collection aggregatorFactories) { @@ -109,10 +107,6 @@ public void incrementSearchQueryAggregationCounters(Collection, Counter> queryHandlers; public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) { this(searchQueryCounters, 0); @@ -60,49 +27,10 @@ public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) { private SearchQueryCategorizingVisitor(SearchQueryCounters counters, int level) { this.searchQueryCounters = counters; this.level = level; - this.queryHandlers = initializeQueryHandlers(); - } - - private Map, Counter> initializeQueryHandlers() { - Map, Counter> handlers = new HashMap<>(); - - handlers.put(BoolQueryBuilder.class, searchQueryCounters.boolCounter); - handlers.put(FunctionScoreQueryBuilder.class, searchQueryCounters.functionScoreCounter); - handlers.put(MatchQueryBuilder.class, searchQueryCounters.matchCounter); - handlers.put(MatchPhraseQueryBuilder.class, searchQueryCounters.matchPhrasePrefixCounter); - handlers.put(MultiMatchQueryBuilder.class, searchQueryCounters.multiMatchCounter); - handlers.put(QueryStringQueryBuilder.class, searchQueryCounters.queryStringQueryCounter); - handlers.put(RangeQueryBuilder.class, searchQueryCounters.rangeCounter); - handlers.put(RegexpQueryBuilder.class, searchQueryCounters.regexCounter); - handlers.put(TermQueryBuilder.class, searchQueryCounters.termCounter); - handlers.put(WildcardQueryBuilder.class, searchQueryCounters.wildcardCounter); - handlers.put(BoostingQueryBuilder.class, searchQueryCounters.boostCounter); - handlers.put(ConstantScoreQueryBuilder.class, searchQueryCounters.constantScoreCounter); - handlers.put(DisMaxQueryBuilder.class, searchQueryCounters.disMaxCounter); - handlers.put(DistanceFeatureQueryBuilder.class, searchQueryCounters.distanceFeatureCounter); - handlers.put(ExistsQueryBuilder.class, searchQueryCounters.existsCounter); - handlers.put(FieldMaskingSpanQueryBuilder.class, searchQueryCounters.fieldMaskingSpanCounter); - handlers.put(FuzzyQueryBuilder.class, searchQueryCounters.fuzzyCounter); - handlers.put(GeoBoundingBoxQueryBuilder.class, searchQueryCounters.geoBoundingBoxCounter); - handlers.put(GeoDistanceQueryBuilder.class, searchQueryCounters.geoDistanceCounter); - handlers.put(GeoPolygonQueryBuilder.class, searchQueryCounters.geoPolygonCounter); - handlers.put(GeoShapeQueryBuilder.class, searchQueryCounters.geoShapeCounter); - handlers.put(IntervalQueryBuilder.class, searchQueryCounters.intervalCounter); - handlers.put(MatchAllQueryBuilder.class, searchQueryCounters.matchallCounter); - handlers.put(PrefixQueryBuilder.class, searchQueryCounters.prefixCounter); - handlers.put(ScriptQueryBuilder.class, searchQueryCounters.scriptCounter); - handlers.put(SimpleQueryStringBuilder.class, searchQueryCounters.simpleQueryStringCounter); - - return handlers; } public void accept(QueryBuilder qb) { - Counter counter = queryHandlers.get(qb.getClass()); - if (counter != null) { - counter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } else { - searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); - } + searchQueryCounters.incrementCounter(qb, level); } public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java index fbbbeb52786e3..bbb883809b41b 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java +++ b/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java @@ -8,20 +8,52 @@ package org.opensearch.action.search; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.BoostingQueryBuilder; +import org.opensearch.index.query.ConstantScoreQueryBuilder; +import org.opensearch.index.query.DisMaxQueryBuilder; +import org.opensearch.index.query.DistanceFeatureQueryBuilder; +import org.opensearch.index.query.ExistsQueryBuilder; +import org.opensearch.index.query.FieldMaskingSpanQueryBuilder; +import org.opensearch.index.query.FuzzyQueryBuilder; +import org.opensearch.index.query.GeoBoundingBoxQueryBuilder; +import org.opensearch.index.query.GeoDistanceQueryBuilder; +import org.opensearch.index.query.GeoPolygonQueryBuilder; +import org.opensearch.index.query.GeoShapeQueryBuilder; +import org.opensearch.index.query.IntervalQueryBuilder; +import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.MatchPhraseQueryBuilder; +import org.opensearch.index.query.MatchQueryBuilder; +import org.opensearch.index.query.MultiMatchQueryBuilder; +import org.opensearch.index.query.PrefixQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryStringQueryBuilder; +import org.opensearch.index.query.RangeQueryBuilder; +import org.opensearch.index.query.RegexpQueryBuilder; +import org.opensearch.index.query.ScriptQueryBuilder; +import org.opensearch.index.query.SimpleQueryStringBuilder; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.index.query.WildcardQueryBuilder; +import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder; import org.opensearch.telemetry.metrics.Counter; import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.metrics.tags.Tags; + +import java.util.HashMap; +import java.util.Map; /** * Class contains all the Counters related to search query types. */ final class SearchQueryCounters { + private static final String LEVEL_TAG = "level"; private static final String UNIT = "1"; private final MetricsRegistry metricsRegistry; // Counters related to Query types public final Counter aggCounter; public final Counter boolCounter; - public final Counter boostCounter; + public final Counter boostingCounter; public final Counter constantScoreCounter; public final Counter disMaxCounter; public final Counter distanceFeatureCounter; @@ -40,9 +72,9 @@ final class SearchQueryCounters { public final Counter multiMatchCounter; public final Counter otherQueryCounter; public final Counter prefixCounter; - public final Counter queryStringQueryCounter; + public final Counter queryStringCounter; public final Counter rangeCounter; - public final Counter regexCounter; + public final Counter regexpCounter; public final Counter scriptCounter; public final Counter simpleQueryStringCounter; public final Counter sortCounter; @@ -51,6 +83,7 @@ final class SearchQueryCounters { public final Counter totalCounter; public final Counter wildcardCounter; public final Counter numberOfInputFieldsCounter; + private final Map, Counter> queryHandlers; public SearchQueryCounters(MetricsRegistry metricsRegistry) { this.metricsRegistry = metricsRegistry; @@ -64,9 +97,9 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) { "Counter for the number of top level and nested bool search queries", UNIT ); - this.boostCounter = metricsRegistry.createCounter( - "search.query.type.boast.count", - "Counter for the number of top level and nested boast search queries", + this.boostingCounter = metricsRegistry.createCounter( + "search.query.type.boost.count", + "Counter for the number of top level and nested boost search queries", UNIT ); this.constantScoreCounter = metricsRegistry.createCounter( @@ -75,12 +108,12 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) { UNIT ); this.disMaxCounter = metricsRegistry.createCounter( - "search.query.type.disMax.count", + "search.query.type.dismax.count", "Counter for the number of top level and nested disjuntion max search queries", UNIT ); this.distanceFeatureCounter = metricsRegistry.createCounter( - "search.query.type.distanceFeature.count", + "search.query.type.distancefeature.count", "Counter for the number of top level and nested distance feature search queries", UNIT ); @@ -159,7 +192,7 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) { "Counter for the number of top level and nested search queries that match prefix queries", UNIT ); - this.queryStringQueryCounter = metricsRegistry.createCounter( + this.queryStringCounter = metricsRegistry.createCounter( "search.query.type.querystringquery.count", "Counter for the number of top level and nested queryStringQuery search queries", UNIT @@ -169,7 +202,7 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) { "Counter for the number of top level and nested range search queries", UNIT ); - this.regexCounter = metricsRegistry.createCounter( + this.regexpCounter = metricsRegistry.createCounter( "search.query.type.regex.count", "Counter for the number of top level and nested regex search queries", UNIT @@ -214,5 +247,46 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) { "Counter for the number of input fields in the search queries", UNIT ); + this.queryHandlers = new HashMap<>(); + initializeQueryHandlers(); + } + + public void incrementCounter(QueryBuilder queryBuilder, int level) { + Counter counter = queryHandlers.get(queryBuilder.getClass()); + if (counter != null) { + counter.add(1, Tags.create().addTag(LEVEL_TAG, level)); + } else { + otherQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level)); + } + } + + private void initializeQueryHandlers() { + + queryHandlers.put(BoolQueryBuilder.class, boolCounter); + queryHandlers.put(FunctionScoreQueryBuilder.class, functionScoreCounter); + queryHandlers.put(MatchQueryBuilder.class, matchCounter); + queryHandlers.put(MatchPhraseQueryBuilder.class, matchPhrasePrefixCounter); + queryHandlers.put(MultiMatchQueryBuilder.class, multiMatchCounter); + queryHandlers.put(QueryStringQueryBuilder.class, queryStringCounter); + queryHandlers.put(RangeQueryBuilder.class, rangeCounter); + queryHandlers.put(RegexpQueryBuilder.class, regexpCounter); + queryHandlers.put(TermQueryBuilder.class, termCounter); + queryHandlers.put(WildcardQueryBuilder.class, wildcardCounter); + queryHandlers.put(BoostingQueryBuilder.class, boostingCounter); + queryHandlers.put(ConstantScoreQueryBuilder.class, constantScoreCounter); + queryHandlers.put(DisMaxQueryBuilder.class, disMaxCounter); + queryHandlers.put(DistanceFeatureQueryBuilder.class, distanceFeatureCounter); + queryHandlers.put(ExistsQueryBuilder.class, existsCounter); + queryHandlers.put(FieldMaskingSpanQueryBuilder.class, fieldMaskingSpanCounter); + queryHandlers.put(FuzzyQueryBuilder.class, fuzzyCounter); + queryHandlers.put(GeoBoundingBoxQueryBuilder.class, geoBoundingBoxCounter); + queryHandlers.put(GeoDistanceQueryBuilder.class, geoDistanceCounter); + queryHandlers.put(GeoPolygonQueryBuilder.class, geoPolygonCounter); + queryHandlers.put(GeoShapeQueryBuilder.class, geoShapeCounter); + queryHandlers.put(IntervalQueryBuilder.class, intervalCounter); + queryHandlers.put(MatchAllQueryBuilder.class, matchallCounter); + queryHandlers.put(PrefixQueryBuilder.class, prefixCounter); + queryHandlers.put(ScriptQueryBuilder.class, scriptCounter); + queryHandlers.put(SimpleQueryStringBuilder.class, simpleQueryStringCounter); } } diff --git a/server/src/test/java/org/opensearch/action/search/SearchQueryCategorizerTests.java b/server/src/test/java/org/opensearch/action/search/SearchQueryCategorizerTests.java index a925395def646..17fa124890158 100644 --- a/server/src/test/java/org/opensearch/action/search/SearchQueryCategorizerTests.java +++ b/server/src/test/java/org/opensearch/action/search/SearchQueryCategorizerTests.java @@ -164,7 +164,7 @@ public void testQueryStringQuery() { searchQueryCategorizer.categorize(sourceBuilder); - verify(searchQueryCategorizer.searchQueryCounters.queryStringQueryCounter).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.searchQueryCounters.queryStringCounter).add(eq(1.0d), any(Tags.class)); } public void testRangeQuery() { @@ -185,7 +185,7 @@ public void testRegexQuery() { searchQueryCategorizer.categorize(sourceBuilder); - verify(searchQueryCategorizer.searchQueryCounters.regexCounter).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.searchQueryCounters.regexpCounter).add(eq(1.0d), any(Tags.class)); } public void testSortQuery() { @@ -237,7 +237,7 @@ public void testComplexQuery() { verify(searchQueryCategorizer.searchQueryCounters.termCounter).add(eq(1.0d), any(Tags.class)); verify(searchQueryCategorizer.searchQueryCounters.matchCounter).add(eq(1.0d), any(Tags.class)); - verify(searchQueryCategorizer.searchQueryCounters.regexCounter).add(eq(1.0d), any(Tags.class)); + verify(searchQueryCategorizer.searchQueryCounters.regexpCounter).add(eq(1.0d), any(Tags.class)); verify(searchQueryCategorizer.searchQueryCounters.boolCounter).add(eq(1.0d), any(Tags.class)); verify(searchQueryCategorizer.searchQueryCounters.aggCounter).add(eq(1.0d), any(Tags.class)); }