Skip to content

Commit

Permalink
Expand wildcard imports and make implementation more succinct
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 Dec 15, 2023
1 parent def4097 commit 13228a2
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 141 deletions.
13 changes: 12 additions & 1 deletion server/src/main/java/org/opensearch/action/search/SearchQueryAggregationCategorizer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,18 @@
import org.opensearch.search.aggregations.bucket.terms.SignificantTermsAggregationBuilder;
import org.opensearch.search.aggregations.bucket.terms.SignificantTextAggregationBuilder;
import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.opensearch.search.aggregations.metrics.*;
import org.opensearch.search.aggregations.metrics.AvgAggregationBuilder;
import org.opensearch.search.aggregations.metrics.CardinalityAggregationBuilder;
import org.opensearch.search.aggregations.metrics.ExtendedStatsAggregationBuilder;
import org.opensearch.search.aggregations.metrics.GeoCentroidAggregationBuilder;
import org.opensearch.search.aggregations.metrics.MaxAggregationBuilder;
import org.opensearch.search.aggregations.metrics.MedianAbsoluteDeviationAggregationBuilder;
import org.opensearch.search.aggregations.metrics.MinAggregationBuilder;
import org.opensearch.search.aggregations.metrics.PercentileRanksAggregationBuilder;
import org.opensearch.search.aggregations.metrics.PercentilesAggregationBuilder;
import org.opensearch.search.aggregations.metrics.ScriptedMetricAggregationBuilder;
import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
import org.opensearch.search.aggregations.metrics.TopHitsAggregationBuilder;
import org.opensearch.telemetry.metrics.tags.Tags;

import java.util.Collection;
Expand Down
204 changes: 64 additions & 140 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,8 +9,35 @@
package org.opensearch.action.search;

import org.apache.lucene.search.BooleanClause;
import org.opensearch.index.query.*;
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;
Expand All @@ -25,7 +52,7 @@ 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>, Consumer<QueryBuilder>> queryHandlers;
private final Map<Class<? extends QueryBuilder>, Counter> queryHandlers;

public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) {
this(searchQueryCounters, 0);
Expand All @@ -37,152 +64,49 @@ private SearchQueryCategorizingVisitor(SearchQueryCounters counters, int level)
this.queryHandlers = initializeQueryHandlers();
}

private Map<Class<? extends QueryBuilder>, Consumer<QueryBuilder>> initializeQueryHandlers() {
Map<Class<? extends QueryBuilder>, Consumer<QueryBuilder>> handlers = new HashMap<>();

handlers.put(BoolQueryBuilder.class, this::handleBoolQuery);
handlers.put(FunctionScoreQueryBuilder.class, this::handleFunctionScoreQuery);
handlers.put(MatchQueryBuilder.class, this::handleMatchQuery);
handlers.put(MatchPhraseQueryBuilder.class, this::handleMatchPhrasePrefixQuery);
handlers.put(MultiMatchQueryBuilder.class, this::handleMultiMatchQuery);
handlers.put(QueryStringQueryBuilder.class, this::handleQueryStringQuery);
handlers.put(RangeQueryBuilder.class, this::handleRangeQuery);
handlers.put(RegexpQueryBuilder.class, this::handleRegexpQuery);
handlers.put(TermQueryBuilder.class, this::handleTermQuery);
handlers.put(WildcardQueryBuilder.class, this::handleWildcardQuery);
handlers.put(BoostingQueryBuilder.class, this::handleBoostingQuery);
handlers.put(ConstantScoreQueryBuilder.class, this::handleConstantScoreQuery);
handlers.put(DisMaxQueryBuilder.class, this::handleDisMaxQuery);
handlers.put(DistanceFeatureQueryBuilder.class, this::handleDistanceFeatureQuery);
handlers.put(ExistsQueryBuilder.class, this::handleExistsQuery);
handlers.put(FieldMaskingSpanQueryBuilder.class, this::handleFieldMaskingSpanQuery);
handlers.put(FuzzyQueryBuilder.class, this::handleFuzzyQuery);
handlers.put(GeoBoundingBoxQueryBuilder.class, this::handleGeoBoundingBoxQuery);
handlers.put(GeoDistanceQueryBuilder.class, this::handleGeoDistanceQuery);
handlers.put(GeoPolygonQueryBuilder.class, this::handleGeoPolygonQuery);
handlers.put(GeoShapeQueryBuilder.class, this::handleGeoShapeQuery);
handlers.put(IntervalQueryBuilder.class, this::handleIntervalQuery);
handlers.put(MatchAllQueryBuilder.class, this::handleMatchAllQuery);
handlers.put(PrefixQueryBuilder.class, this::handlePrefixQuery);
handlers.put(ScriptQueryBuilder.class, this::handleScriptQuery);
handlers.put(SimpleQueryStringBuilder.class, this::handleSimpleQueryStringQuery);
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) {
queryHandlers.getOrDefault(qb.getClass(), this::handleOtherQuery).accept(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));
}
}

public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
return new SearchQueryCategorizingVisitor(searchQueryCounters, level + 1);
}

private void handleBoolQuery(QueryBuilder qb) {
searchQueryCounters.boolCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleFunctionScoreQuery(QueryBuilder qb) {
searchQueryCounters.functionScoreCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleMatchQuery(QueryBuilder qb) {
searchQueryCounters.matchCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleMatchPhrasePrefixQuery(QueryBuilder qb) {
searchQueryCounters.matchPhrasePrefixCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleMultiMatchQuery(QueryBuilder qb) {
searchQueryCounters.multiMatchCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleQueryStringQuery(QueryBuilder qb) {
searchQueryCounters.queryStringQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleRangeQuery(QueryBuilder qb) {
searchQueryCounters.rangeCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleRegexpQuery(QueryBuilder qb) {
searchQueryCounters.regexCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleTermQuery(QueryBuilder qb) {
searchQueryCounters.termCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleWildcardQuery(QueryBuilder qb) {
searchQueryCounters.wildcardCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleBoostingQuery(QueryBuilder qb) {
searchQueryCounters.boostCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleConstantScoreQuery(QueryBuilder qb) {
searchQueryCounters.constantScoreCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleDisMaxQuery(QueryBuilder qb) {
searchQueryCounters.disMaxCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleDistanceFeatureQuery(QueryBuilder qb) {
searchQueryCounters.distanceFeatureCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleExistsQuery(QueryBuilder qb) {
searchQueryCounters.existsCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleFieldMaskingSpanQuery(QueryBuilder qb) {
searchQueryCounters.fieldMaskingSpanCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleFuzzyQuery(QueryBuilder qb) {
searchQueryCounters.fuzzyCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleGeoBoundingBoxQuery(QueryBuilder qb) {
searchQueryCounters.geoBoundingBoxCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleGeoDistanceQuery(QueryBuilder qb) {
searchQueryCounters.geoDistanceCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleGeoPolygonQuery(QueryBuilder qb) {
searchQueryCounters.geoPolygonCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleGeoShapeQuery(QueryBuilder qb) {
searchQueryCounters.geoShapeCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleIntervalQuery(QueryBuilder qb) {
searchQueryCounters.intervalCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleMatchAllQuery(QueryBuilder qb) {
searchQueryCounters.matchallCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handlePrefixQuery(QueryBuilder qb) {
searchQueryCounters.prefixCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleScriptQuery(QueryBuilder qb) {
searchQueryCounters.scriptCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleSimpleQueryStringQuery(QueryBuilder qb) {
searchQueryCounters.simpleQueryStringCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}

private void handleOtherQuery(QueryBuilder qb) {
searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
}
}

0 comments on commit 13228a2

Please sign in to comment.