Skip to content

Commit 8539426

Browse files
dzane17dzane17
committed
Refactor query shape field data maps
Signed-off-by: David Zane <davizane@amazon.com>
1 parent 290c701 commit 8539426

File tree

2 files changed

+9
-154
lines changed

2 files changed

+9
-154
lines changed

‎src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java

+6-142
Original file line numberDiff line numberDiff line change
@@ -12,60 +12,15 @@
1212
import java.util.Collection;
1313
import java.util.Collections;
1414
import java.util.List;
15-
import java.util.Map;
16-
import java.util.function.Function;
1715
import org.apache.lucene.util.BytesRef;
1816
import org.opensearch.common.hash.MurmurHash3;
1917
import org.opensearch.core.common.io.stream.NamedWriteable;
20-
import org.opensearch.index.query.AbstractGeometryQueryBuilder;
21-
import org.opensearch.index.query.CommonTermsQueryBuilder;
22-
import org.opensearch.index.query.ExistsQueryBuilder;
23-
import org.opensearch.index.query.FieldMaskingSpanQueryBuilder;
24-
import org.opensearch.index.query.FuzzyQueryBuilder;
25-
import org.opensearch.index.query.GeoBoundingBoxQueryBuilder;
26-
import org.opensearch.index.query.GeoDistanceQueryBuilder;
27-
import org.opensearch.index.query.GeoPolygonQueryBuilder;
28-
import org.opensearch.index.query.MatchBoolPrefixQueryBuilder;
29-
import org.opensearch.index.query.MatchPhrasePrefixQueryBuilder;
30-
import org.opensearch.index.query.MatchPhraseQueryBuilder;
31-
import org.opensearch.index.query.MatchQueryBuilder;
32-
import org.opensearch.index.query.MultiTermQueryBuilder;
33-
import org.opensearch.index.query.PrefixQueryBuilder;
3418
import org.opensearch.index.query.QueryBuilder;
35-
import org.opensearch.index.query.RangeQueryBuilder;
36-
import org.opensearch.index.query.RegexpQueryBuilder;
37-
import org.opensearch.index.query.SpanNearQueryBuilder;
38-
import org.opensearch.index.query.SpanTermQueryBuilder;
39-
import org.opensearch.index.query.TermQueryBuilder;
40-
import org.opensearch.index.query.TermsQueryBuilder;
41-
import org.opensearch.index.query.WildcardQueryBuilder;
19+
import org.opensearch.index.query.WithFieldName;
4220
import org.opensearch.search.aggregations.AggregationBuilder;
4321
import org.opensearch.search.aggregations.AggregatorFactories;
4422
import org.opensearch.search.aggregations.PipelineAggregationBuilder;
45-
import org.opensearch.search.aggregations.bucket.histogram.AutoDateHistogramAggregationBuilder;
46-
import org.opensearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder;
47-
import org.opensearch.search.aggregations.bucket.histogram.HistogramAggregationBuilder;
48-
import org.opensearch.search.aggregations.bucket.histogram.VariableWidthHistogramAggregationBuilder;
49-
import org.opensearch.search.aggregations.bucket.missing.MissingAggregationBuilder;
50-
import org.opensearch.search.aggregations.bucket.range.AbstractRangeBuilder;
51-
import org.opensearch.search.aggregations.bucket.range.GeoDistanceAggregationBuilder;
52-
import org.opensearch.search.aggregations.bucket.range.IpRangeAggregationBuilder;
53-
import org.opensearch.search.aggregations.bucket.sampler.DiversifiedAggregationBuilder;
54-
import org.opensearch.search.aggregations.bucket.terms.RareTermsAggregationBuilder;
55-
import org.opensearch.search.aggregations.bucket.terms.SignificantTermsAggregationBuilder;
56-
import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
57-
import org.opensearch.search.aggregations.metrics.AvgAggregationBuilder;
58-
import org.opensearch.search.aggregations.metrics.CardinalityAggregationBuilder;
59-
import org.opensearch.search.aggregations.metrics.ExtendedStatsAggregationBuilder;
60-
import org.opensearch.search.aggregations.metrics.GeoCentroidAggregationBuilder;
61-
import org.opensearch.search.aggregations.metrics.MaxAggregationBuilder;
62-
import org.opensearch.search.aggregations.metrics.MinAggregationBuilder;
63-
import org.opensearch.search.aggregations.metrics.StatsAggregationBuilder;
64-
import org.opensearch.search.aggregations.metrics.SumAggregationBuilder;
65-
import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
66-
import org.opensearch.search.aggregations.support.ValuesSourceAggregationBuilder;
6723
import org.opensearch.search.builder.SearchSourceBuilder;
68-
import org.opensearch.search.sort.FieldSortBuilder;
6924
import org.opensearch.search.sort.SortBuilder;
7025

7126
/**
@@ -74,9 +29,6 @@
7429
public class QueryShapeGenerator {
7530
static final String EMPTY_STRING = "";
7631
static final String ONE_SPACE_INDENT = " ";
77-
static final Map<Class<?>, List<Function<Object, String>>> QUERY_FIELD_DATA_MAP = FieldDataMapHelper.getQueryFieldDataMap();
78-
static final Map<Class<?>, List<Function<Object, String>>> AGG_FIELD_DATA_MAP = FieldDataMapHelper.getAggFieldDataMap();
79-
static final Map<Class<?>, List<Function<Object, String>>> SORT_FIELD_DATA_MAP = FieldDataMapHelper.getSortFieldDataMap();
8032

8133
/**
8234
* Method to get query shape hash code given a source
@@ -161,7 +113,7 @@ static StringBuilder recursiveAggregationShapeBuilder(
161113
StringBuilder stringBuilder = new StringBuilder();
162114
stringBuilder.append(baseIndent).append(ONE_SPACE_INDENT.repeat(2)).append(aggBuilder.getType());
163115
if (showFields) {
164-
stringBuilder.append(buildFieldDataString(AGG_FIELD_DATA_MAP.get(aggBuilder.getClass()), aggBuilder));
116+
stringBuilder.append(buildFieldDataString(aggBuilder));
165117
}
166118
stringBuilder.append("\n");
167119

@@ -227,7 +179,7 @@ static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showF
227179
StringBuilder stringBuilder = new StringBuilder();
228180
stringBuilder.append(ONE_SPACE_INDENT.repeat(2)).append(sortBuilder.order());
229181
if (showFields) {
230-
stringBuilder.append(buildFieldDataString(SORT_FIELD_DATA_MAP.get(sortBuilder.getClass()), sortBuilder));
182+
stringBuilder.append(buildFieldDataString(sortBuilder));
231183
}
232184
shapeStrings.add(stringBuilder.toString());
233185
}
@@ -243,99 +195,11 @@ static String buildSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showF
243195
* @return String: comma separated list with leading space in square brackets
244196
* Ex: " [my_field, width:5]"
245197
*/
246-
static String buildFieldDataString(List<Function<Object, String>> methods, NamedWriteable builder) {
198+
static String buildFieldDataString(NamedWriteable builder) {
247199
List<String> fieldDataList = new ArrayList<>();
248-
if (methods != null) {
249-
for (Function<Object, String> lambda : methods) {
250-
fieldDataList.add(lambda.apply(builder));
251-
}
200+
if (builder instanceof WithFieldName) {
201+
fieldDataList.add(((WithFieldName) builder).fieldName());
252202
}
253203
return " [" + String.join(", ", fieldDataList) + "]";
254204
}
255-
256-
/**
257-
* Helper class to create static field data maps
258-
*/
259-
private static class FieldDataMapHelper {
260-
261-
// Helper method to create map entries
262-
private static <T> Map.Entry<Class<?>, List<Function<Object, String>>> createEntry(Class<T> clazz, Function<T, String> extractor) {
263-
return Map.entry(clazz, List.of(obj -> extractor.apply(clazz.cast(obj))));
264-
}
265-
266-
/**
267-
* Returns a map where the keys are query builders, and the values are lists of
268-
* functions that extract field values from instances of these classes.
269-
*
270-
* @return a map with class types as keys and lists of field extraction functions as values.
271-
*/
272-
private static Map<Class<?>, List<Function<Object, String>>> getQueryFieldDataMap() {
273-
return Map.ofEntries(
274-
createEntry(AbstractGeometryQueryBuilder.class, AbstractGeometryQueryBuilder::fieldName),
275-
createEntry(CommonTermsQueryBuilder.class, CommonTermsQueryBuilder::fieldName),
276-
createEntry(ExistsQueryBuilder.class, ExistsQueryBuilder::fieldName),
277-
createEntry(FieldMaskingSpanQueryBuilder.class, FieldMaskingSpanQueryBuilder::fieldName),
278-
createEntry(FuzzyQueryBuilder.class, FuzzyQueryBuilder::fieldName),
279-
createEntry(GeoBoundingBoxQueryBuilder.class, GeoBoundingBoxQueryBuilder::fieldName),
280-
createEntry(GeoDistanceQueryBuilder.class, GeoDistanceQueryBuilder::fieldName),
281-
createEntry(GeoPolygonQueryBuilder.class, GeoPolygonQueryBuilder::fieldName),
282-
createEntry(MatchBoolPrefixQueryBuilder.class, MatchBoolPrefixQueryBuilder::fieldName),
283-
createEntry(MatchQueryBuilder.class, MatchQueryBuilder::fieldName),
284-
createEntry(MatchPhraseQueryBuilder.class, MatchPhraseQueryBuilder::fieldName),
285-
createEntry(MatchPhrasePrefixQueryBuilder.class, MatchPhrasePrefixQueryBuilder::fieldName),
286-
createEntry(MultiTermQueryBuilder.class, MultiTermQueryBuilder::fieldName),
287-
createEntry(PrefixQueryBuilder.class, PrefixQueryBuilder::fieldName),
288-
createEntry(RangeQueryBuilder.class, RangeQueryBuilder::fieldName),
289-
createEntry(RegexpQueryBuilder.class, RegexpQueryBuilder::fieldName),
290-
createEntry(SpanNearQueryBuilder.SpanGapQueryBuilder.class, SpanNearQueryBuilder.SpanGapQueryBuilder::fieldName),
291-
createEntry(SpanTermQueryBuilder.class, SpanTermQueryBuilder::fieldName),
292-
createEntry(TermQueryBuilder.class, TermQueryBuilder::fieldName),
293-
createEntry(TermsQueryBuilder.class, TermsQueryBuilder::fieldName),
294-
createEntry(WildcardQueryBuilder.class, WildcardQueryBuilder::fieldName)
295-
);
296-
}
297-
298-
/**
299-
* Returns a map where the keys are aggregation builders, and the values are lists of
300-
* functions that extract field values from instances of these classes.
301-
*
302-
* @return a map with class types as keys and lists of field extraction functions as values.
303-
*/
304-
private static Map<Class<?>, List<Function<Object, String>>> getAggFieldDataMap() {
305-
return Map.ofEntries(
306-
createEntry(IpRangeAggregationBuilder.class, IpRangeAggregationBuilder::field),
307-
createEntry(AutoDateHistogramAggregationBuilder.class, AutoDateHistogramAggregationBuilder::field),
308-
createEntry(DateHistogramAggregationBuilder.class, DateHistogramAggregationBuilder::field),
309-
createEntry(HistogramAggregationBuilder.class, HistogramAggregationBuilder::field),
310-
createEntry(VariableWidthHistogramAggregationBuilder.class, VariableWidthHistogramAggregationBuilder::field),
311-
createEntry(MissingAggregationBuilder.class, MissingAggregationBuilder::field),
312-
createEntry(AbstractRangeBuilder.class, AbstractRangeBuilder::field),
313-
createEntry(GeoDistanceAggregationBuilder.class, GeoDistanceAggregationBuilder::field),
314-
createEntry(DiversifiedAggregationBuilder.class, DiversifiedAggregationBuilder::field),
315-
createEntry(RareTermsAggregationBuilder.class, RareTermsAggregationBuilder::field),
316-
createEntry(SignificantTermsAggregationBuilder.class, SignificantTermsAggregationBuilder::field),
317-
createEntry(TermsAggregationBuilder.class, TermsAggregationBuilder::field),
318-
createEntry(AvgAggregationBuilder.class, AvgAggregationBuilder::field),
319-
createEntry(CardinalityAggregationBuilder.class, CardinalityAggregationBuilder::field),
320-
createEntry(ExtendedStatsAggregationBuilder.class, ExtendedStatsAggregationBuilder::field),
321-
createEntry(GeoCentroidAggregationBuilder.class, GeoCentroidAggregationBuilder::field),
322-
createEntry(MaxAggregationBuilder.class, MaxAggregationBuilder::field),
323-
createEntry(MinAggregationBuilder.class, MinAggregationBuilder::field),
324-
createEntry(StatsAggregationBuilder.class, StatsAggregationBuilder::field),
325-
createEntry(SumAggregationBuilder.class, SumAggregationBuilder::field),
326-
createEntry(ValueCountAggregationBuilder.class, ValueCountAggregationBuilder::field),
327-
createEntry(ValuesSourceAggregationBuilder.class, ValuesSourceAggregationBuilder::field)
328-
);
329-
}
330-
331-
/**
332-
* Returns a map where the keys are sort builders, and the values are lists of
333-
* functions that extract field values from instances of these classes.
334-
*
335-
* @return a map with class types as keys and lists of field extraction functions as values.
336-
*/
337-
private static Map<Class<?>, List<Function<Object, String>>> getSortFieldDataMap() {
338-
return Map.ofEntries(createEntry(FieldSortBuilder.class, FieldSortBuilder::getFieldName));
339-
}
340-
}
341205
}

‎src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeVisitor.java

+3-12
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,13 @@
99
package org.opensearch.plugin.insights.core.service.categorizer;
1010

1111
import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.ONE_SPACE_INDENT;
12-
import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.QUERY_FIELD_DATA_MAP;
12+
import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.buildFieldDataString;
1313

1414
import java.util.ArrayList;
1515
import java.util.EnumMap;
1616
import java.util.List;
1717
import java.util.Locale;
1818
import java.util.Map;
19-
import java.util.function.Function;
2019
import org.apache.lucene.search.BooleanClause;
2120
import org.opensearch.common.SetOnce;
2221
import org.opensearch.index.query.QueryBuilder;
@@ -33,15 +32,7 @@ public final class QueryShapeVisitor implements QueryBuilderVisitor {
3332
@Override
3433
public void accept(QueryBuilder queryBuilder) {
3534
queryType.set(queryBuilder.getName());
36-
37-
List<String> fieldDataList = new ArrayList<>();
38-
List<Function<Object, String>> methods = QUERY_FIELD_DATA_MAP.get(queryBuilder.getClass());
39-
if (methods != null) {
40-
for (Function<Object, String> lambda : methods) {
41-
fieldDataList.add(lambda.apply(queryBuilder));
42-
}
43-
}
44-
fieldData.set(String.join(", ", fieldDataList));
35+
fieldData.set(buildFieldDataString(queryBuilder));
4536
}
4637

4738
@Override
@@ -101,7 +92,7 @@ public String toJson() {
10192
public String prettyPrintTree(String indent, Boolean showFields) {
10293
StringBuilder outputBuilder = new StringBuilder(indent).append(queryType.get());
10394
if (showFields) {
104-
outputBuilder.append(" [").append(fieldData.get()).append("]");
95+
outputBuilder.append(fieldData.get());
10596
}
10697
outputBuilder.append("\n");
10798
for (Map.Entry<BooleanClause.Occur, List<QueryShapeVisitor>> entry : childVisitors.entrySet()) {

0 commit comments

Comments
 (0)