From 030b5ef30a758052e582e62e8a6325a703c3b81e Mon Sep 17 00:00:00 2001 From: David Zane Date: Mon, 16 Sep 2024 17:30:50 -0700 Subject: [PATCH] Refactor query shape field data maps Signed-off-by: David Zane --- .../categorizer/QueryShapeGenerator.java | 108 ++---------------- .../categorizer/QueryShapeVisitor.java | 14 +-- 2 files changed, 15 insertions(+), 107 deletions(-) diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java index 5664f3c..7b28f3a 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java @@ -17,53 +17,11 @@ import org.apache.lucene.util.BytesRef; import org.opensearch.common.hash.MurmurHash3; import org.opensearch.core.common.io.stream.NamedWriteable; -import org.opensearch.index.query.AbstractGeometryQueryBuilder; -import org.opensearch.index.query.CommonTermsQueryBuilder; -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.MatchBoolPrefixQueryBuilder; -import org.opensearch.index.query.MatchPhrasePrefixQueryBuilder; -import org.opensearch.index.query.MatchPhraseQueryBuilder; -import org.opensearch.index.query.MatchQueryBuilder; -import org.opensearch.index.query.MultiTermQueryBuilder; -import org.opensearch.index.query.PrefixQueryBuilder; import org.opensearch.index.query.QueryBuilder; -import org.opensearch.index.query.RangeQueryBuilder; -import org.opensearch.index.query.RegexpQueryBuilder; -import org.opensearch.index.query.SpanNearQueryBuilder; -import org.opensearch.index.query.SpanTermQueryBuilder; -import org.opensearch.index.query.TermQueryBuilder; -import org.opensearch.index.query.TermsQueryBuilder; -import org.opensearch.index.query.WildcardQueryBuilder; +import org.opensearch.index.query.WithFieldName; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.AggregatorFactories; import org.opensearch.search.aggregations.PipelineAggregationBuilder; -import org.opensearch.search.aggregations.bucket.histogram.AutoDateHistogramAggregationBuilder; -import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder; -import org.opensearch.search.aggregations.bucket.histogram.HistogramAggregationBuilder; -import org.opensearch.search.aggregations.bucket.histogram.VariableWidthHistogramAggregationBuilder; -import org.opensearch.search.aggregations.bucket.missing.MissingAggregationBuilder; -import org.opensearch.search.aggregations.bucket.range.AbstractRangeBuilder; -import org.opensearch.search.aggregations.bucket.range.GeoDistanceAggregationBuilder; -import org.opensearch.search.aggregations.bucket.range.IpRangeAggregationBuilder; -import org.opensearch.search.aggregations.bucket.sampler.DiversifiedAggregationBuilder; -import org.opensearch.search.aggregations.bucket.terms.RareTermsAggregationBuilder; -import org.opensearch.search.aggregations.bucket.terms.SignificantTermsAggregationBuilder; -import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder; -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.MinAggregationBuilder; -import org.opensearch.search.aggregations.metrics.StatsAggregationBuilder; -import org.opensearch.search.aggregations.metrics.SumAggregationBuilder; -import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder; -import org.opensearch.search.aggregations.support.ValuesSourceAggregationBuilder; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.sort.FieldSortBuilder; import org.opensearch.search.sort.SortBuilder; @@ -161,7 +119,7 @@ static StringBuilder recursiveAggregationShapeBuilder( StringBuilder stringBuilder = new StringBuilder(); stringBuilder.append(baseIndent).append(ONE_SPACE_INDENT.repeat(2)).append(aggBuilder.getType()); if (showFields) { - stringBuilder.append(buildFieldDataString(AGG_FIELD_DATA_MAP.get(aggBuilder.getClass()), aggBuilder)); + stringBuilder.append(buildFieldDataString(AGG_FIELD_DATA_MAP, aggBuilder)); } stringBuilder.append("\n"); @@ -227,7 +185,7 @@ static String buildSortShape(List> sortBuilderList, Boolean showF StringBuilder stringBuilder = new StringBuilder(); stringBuilder.append(ONE_SPACE_INDENT.repeat(2)).append(sortBuilder.order()); if (showFields) { - stringBuilder.append(buildFieldDataString(SORT_FIELD_DATA_MAP.get(sortBuilder.getClass()), sortBuilder)); + stringBuilder.append(buildFieldDataString(SORT_FIELD_DATA_MAP, sortBuilder)); } shapeStrings.add(stringBuilder.toString()); } @@ -243,11 +201,14 @@ static String buildSortShape(List> sortBuilderList, Boolean showF * @return String: comma separated list with leading space in square brackets * Ex: " [my_field, width:5]" */ - static String buildFieldDataString(List> methods, NamedWriteable builder) { + static String buildFieldDataString(Map, List>> fieldMethodsMap, NamedWriteable builder) { List fieldDataList = new ArrayList<>(); - if (methods != null) { - for (Function lambda : methods) { - fieldDataList.add(lambda.apply(builder)); + for (Map.Entry, List>> entry : fieldMethodsMap.entrySet()) { + if (entry.getKey().isInstance(builder)) { + List> methodsList = entry.getValue(); + for (Function lambda : methodsList) { + fieldDataList.add(lambda.apply(builder)); + } } } return " [" + String.join(", ", fieldDataList) + "]"; @@ -270,29 +231,7 @@ private static Map.Entry, List>> createEnt * @return a map with class types as keys and lists of field extraction functions as values. */ private static Map, List>> getQueryFieldDataMap() { - return Map.ofEntries( - createEntry(AbstractGeometryQueryBuilder.class, AbstractGeometryQueryBuilder::fieldName), - createEntry(CommonTermsQueryBuilder.class, CommonTermsQueryBuilder::fieldName), - createEntry(ExistsQueryBuilder.class, ExistsQueryBuilder::fieldName), - createEntry(FieldMaskingSpanQueryBuilder.class, FieldMaskingSpanQueryBuilder::fieldName), - createEntry(FuzzyQueryBuilder.class, FuzzyQueryBuilder::fieldName), - createEntry(GeoBoundingBoxQueryBuilder.class, GeoBoundingBoxQueryBuilder::fieldName), - createEntry(GeoDistanceQueryBuilder.class, GeoDistanceQueryBuilder::fieldName), - createEntry(GeoPolygonQueryBuilder.class, GeoPolygonQueryBuilder::fieldName), - createEntry(MatchBoolPrefixQueryBuilder.class, MatchBoolPrefixQueryBuilder::fieldName), - createEntry(MatchQueryBuilder.class, MatchQueryBuilder::fieldName), - createEntry(MatchPhraseQueryBuilder.class, MatchPhraseQueryBuilder::fieldName), - createEntry(MatchPhrasePrefixQueryBuilder.class, MatchPhrasePrefixQueryBuilder::fieldName), - createEntry(MultiTermQueryBuilder.class, MultiTermQueryBuilder::fieldName), - createEntry(PrefixQueryBuilder.class, PrefixQueryBuilder::fieldName), - createEntry(RangeQueryBuilder.class, RangeQueryBuilder::fieldName), - createEntry(RegexpQueryBuilder.class, RegexpQueryBuilder::fieldName), - createEntry(SpanNearQueryBuilder.SpanGapQueryBuilder.class, SpanNearQueryBuilder.SpanGapQueryBuilder::fieldName), - createEntry(SpanTermQueryBuilder.class, SpanTermQueryBuilder::fieldName), - createEntry(TermQueryBuilder.class, TermQueryBuilder::fieldName), - createEntry(TermsQueryBuilder.class, TermsQueryBuilder::fieldName), - createEntry(WildcardQueryBuilder.class, WildcardQueryBuilder::fieldName) - ); + return Map.ofEntries(createEntry(WithFieldName.class, WithFieldName::fieldName)); } /** @@ -302,30 +241,7 @@ private static Map, List>> getQueryFieldDataMa * @return a map with class types as keys and lists of field extraction functions as values. */ private static Map, List>> getAggFieldDataMap() { - return Map.ofEntries( - createEntry(IpRangeAggregationBuilder.class, IpRangeAggregationBuilder::field), - createEntry(AutoDateHistogramAggregationBuilder.class, AutoDateHistogramAggregationBuilder::field), - createEntry(DateHistogramAggregationBuilder.class, DateHistogramAggregationBuilder::field), - createEntry(HistogramAggregationBuilder.class, HistogramAggregationBuilder::field), - createEntry(VariableWidthHistogramAggregationBuilder.class, VariableWidthHistogramAggregationBuilder::field), - createEntry(MissingAggregationBuilder.class, MissingAggregationBuilder::field), - createEntry(AbstractRangeBuilder.class, AbstractRangeBuilder::field), - createEntry(GeoDistanceAggregationBuilder.class, GeoDistanceAggregationBuilder::field), - createEntry(DiversifiedAggregationBuilder.class, DiversifiedAggregationBuilder::field), - createEntry(RareTermsAggregationBuilder.class, RareTermsAggregationBuilder::field), - createEntry(SignificantTermsAggregationBuilder.class, SignificantTermsAggregationBuilder::field), - createEntry(TermsAggregationBuilder.class, TermsAggregationBuilder::field), - createEntry(AvgAggregationBuilder.class, AvgAggregationBuilder::field), - createEntry(CardinalityAggregationBuilder.class, CardinalityAggregationBuilder::field), - createEntry(ExtendedStatsAggregationBuilder.class, ExtendedStatsAggregationBuilder::field), - createEntry(GeoCentroidAggregationBuilder.class, GeoCentroidAggregationBuilder::field), - createEntry(MaxAggregationBuilder.class, MaxAggregationBuilder::field), - createEntry(MinAggregationBuilder.class, MinAggregationBuilder::field), - createEntry(StatsAggregationBuilder.class, StatsAggregationBuilder::field), - createEntry(SumAggregationBuilder.class, SumAggregationBuilder::field), - createEntry(ValueCountAggregationBuilder.class, ValueCountAggregationBuilder::field), - createEntry(ValuesSourceAggregationBuilder.class, ValuesSourceAggregationBuilder::field) - ); + return Map.ofEntries(createEntry(WithFieldName.class, WithFieldName::fieldName)); } /** diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java index b064200..9c573fe 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java @@ -10,13 +10,13 @@ import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.ONE_SPACE_INDENT; import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.QUERY_FIELD_DATA_MAP; +import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.buildFieldDataString; import java.util.ArrayList; import java.util.EnumMap; import java.util.List; import java.util.Locale; import java.util.Map; -import java.util.function.Function; import org.apache.lucene.search.BooleanClause; import org.opensearch.common.SetOnce; import org.opensearch.index.query.QueryBuilder; @@ -33,15 +33,7 @@ public final class QueryShapeVisitor implements QueryBuilderVisitor { @Override public void accept(QueryBuilder queryBuilder) { queryType.set(queryBuilder.getName()); - - List fieldDataList = new ArrayList<>(); - List> methods = QUERY_FIELD_DATA_MAP.get(queryBuilder.getClass()); - if (methods != null) { - for (Function lambda : methods) { - fieldDataList.add(lambda.apply(queryBuilder)); - } - } - fieldData.set(String.join(", ", fieldDataList)); + fieldData.set(buildFieldDataString(QUERY_FIELD_DATA_MAP, queryBuilder)); } @Override @@ -101,7 +93,7 @@ public String toJson() { public String prettyPrintTree(String indent, Boolean showFields) { StringBuilder outputBuilder = new StringBuilder(indent).append(queryType.get()); if (showFields) { - outputBuilder.append(" [").append(fieldData.get()).append("]"); + outputBuilder.append(fieldData.get()); } outputBuilder.append("\n"); for (Map.Entry> entry : childVisitors.entrySet()) {