Skip to content

Commit

Permalink
Address review comments, refactor code
Browse files Browse the repository at this point in the history 
Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com>
  • Loading branch information
@deshsidd
deshsidd committed Jan 3, 2024
1 parent bdecac6 commit 3f6a003
Show file tree
Hide file tree
Showing 4 changed files with 126 additions and 130 deletions.
82 changes: 38 additions & 44 deletions server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -55,50 +54,49 @@ public class SearchQueryAggregationCategorizer {

private static final String TYPE_TAG = "type";
private final SearchQueryCounters searchQueryCounters;
private final Map<Class<? extends AggregationBuilder>, Function<AggregationBuilder, String>> aggregationHandlers;
private final Map<Class<? extends AggregationBuilder>, String> aggregationHandlers;

public SearchQueryAggregationCategorizer(SearchQueryCounters searchQueryCounters) {
this.searchQueryCounters = searchQueryCounters;
this.aggregationHandlers = initializeAggregationHandlers();
}
this.aggregationHandlers = new HashMap<>();

private Map<Class<? extends AggregationBuilder>, Function<AggregationBuilder, String>> initializeAggregationHandlers() {
Map<Class<? extends AggregationBuilder>, Function<AggregationBuilder, String>> 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<AggregationBuilder> aggregatorFactories) {
Expand All @@ -109,10 +107,6 @@ public void incrementSearchQueryAggregationCounters(Collection<AggregationBuilde
}

private String getAggregationType(AggregationBuilder aggregationBuilder) {
return aggregationHandlers.getOrDefault(aggregationBuilder.getClass(), this::handleUnknownAggregation).apply(aggregationBuilder);
}

private String handleUnknownAggregation(AggregationBuilder aggregationBuilder) {
return "other";
return aggregationHandlers.getOrDefault(aggregationBuilder.getClass(), "other");
}
}
74 changes: 1 addition & 73 deletions server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,49 +9,16 @@
package org.opensearch.action.search;

import org.apache.lucene.search.BooleanClause;
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.QueryBuilderVisitor;
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.tags.Tags;

import java.util.HashMap;
import java.util.Map;

/**
* Class to visit the query builder tree and also track the level information.
* Increments the counters related to Search Query type.
*/
final class SearchQueryCategorizingVisitor implements QueryBuilderVisitor {
private static final String LEVEL_TAG = "level";
private final int level;
private final SearchQueryCounters searchQueryCounters;
private final Map<Class<? extends QueryBuilder>, Counter> queryHandlers;

public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) {
this(searchQueryCounters, 0);
Expand All @@ -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<Class<? extends QueryBuilder>, Counter> initializeQueryHandlers() {
Map<Class<? extends QueryBuilder>, 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) {
Expand Down
Loading

0 comments on commit 3f6a003

Please sign in to comment.