Query field data

Signed-off-by: David Zane <[email protected]>
dzane17 · Jul 25, 2024 · 842d628 · 842d628
1 parent 3bee0cf
commit 842d628
Show file tree

Hide file tree

Showing 3 changed files with 128 additions and 23 deletions.
diff --git a/...ain/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java b/...ain/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeGenerator.java
@@ -8,7 +8,27 @@
 
 package org.opensearch.plugin.insights.core.service.categorizer;
 
+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.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.search.aggregations.AggregationBuilder;
 import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.aggregations.PipelineAggregationBuilder;
@@ -21,25 +41,59 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
 
 public class QueryShapeGenerator {
     static final String TWO_SPACE_INDENT = "  ";
+    static final Map<Class<?>, List<Function<Object, String>>> QUERY_FIELD_DATA_MAP = Map.ofEntries(
+        Map.entry(AbstractGeometryQueryBuilder.class, List.of(obj -> ((AbstractGeometryQueryBuilder) obj).fieldName())),
+        Map.entry(CommonTermsQueryBuilder.class, List.of(obj -> ((CommonTermsQueryBuilder) obj).fieldName())),
+        Map.entry(org.opensearch.index.query.ExistsQueryBuilder.class, List.of(obj -> ((ExistsQueryBuilder) obj).fieldName())),
+        Map.entry(
+            org.opensearch.index.query.FieldMaskingSpanQueryBuilder.class,
+            List.of(obj -> ((FieldMaskingSpanQueryBuilder) obj).fieldName())
+        ),
+        Map.entry(FuzzyQueryBuilder.class, List.of(obj -> ((FuzzyQueryBuilder) obj).fieldName())),
+        Map.entry(
+            org.opensearch.index.query.GeoBoundingBoxQueryBuilder.class,
+            List.of(obj -> ((org.opensearch.index.query.GeoBoundingBoxQueryBuilder) obj).fieldName())
+        ),
+        Map.entry(org.opensearch.index.query.GeoDistanceQueryBuilder.class, List.of(obj -> ((GeoDistanceQueryBuilder) obj).fieldName())),
+        Map.entry(GeoPolygonQueryBuilder.class, List.of(obj -> ((GeoPolygonQueryBuilder) obj).fieldName())),
+        Map.entry(MatchBoolPrefixQueryBuilder.class, List.of(obj -> ((MatchBoolPrefixQueryBuilder) obj).fieldName())),
+        Map.entry(MatchQueryBuilder.class, List.of(obj -> ((MatchQueryBuilder) obj).fieldName())),
+        Map.entry(org.opensearch.index.query.MatchPhraseQueryBuilder.class, List.of(obj -> ((MatchPhraseQueryBuilder) obj).fieldName())),
+        Map.entry(MatchPhrasePrefixQueryBuilder.class, List.of(obj -> ((MatchPhrasePrefixQueryBuilder) obj).fieldName())),
+        Map.entry(MultiTermQueryBuilder.class, List.of(obj -> ((MultiTermQueryBuilder) obj).fieldName())),
+        Map.entry(PrefixQueryBuilder.class, List.of(obj -> ((PrefixQueryBuilder) obj).fieldName())),
+        Map.entry(RangeQueryBuilder.class, List.of(obj -> ((RangeQueryBuilder) obj).fieldName())),
+        Map.entry(RegexpQueryBuilder.class, List.of(obj -> ((RegexpQueryBuilder) obj).fieldName())),
+        Map.entry(
+            SpanNearQueryBuilder.SpanGapQueryBuilder.class,
+            List.of(obj -> ((SpanNearQueryBuilder.SpanGapQueryBuilder) obj).fieldName())
+        ),
+        Map.entry(SpanTermQueryBuilder.class, List.of(obj -> ((SpanTermQueryBuilder) obj).fieldName())),
+        Map.entry(TermQueryBuilder.class, List.of(obj -> ((TermQueryBuilder) obj).fieldName())),
+        Map.entry(TermsQueryBuilder.class, List.of(obj -> ((TermsQueryBuilder) obj).fieldName())),
+        Map.entry(WildcardQueryBuilder.class, List.of(obj -> ((WildcardQueryBuilder) obj).fieldName()))
+    );
 
     public static String buildShape(SearchSourceBuilder source, Boolean showFields) {
         StringBuilder shape = new StringBuilder();
-        shape.append(buildQueryShape(source.query()));
+        shape.append(buildQueryShape(source.query(), showFields));
         shape.append(buildAggregationShape(source.aggregations(), showFields));
         shape.append(buildSortShape(source.sorts(), showFields));
         return shape.toString();
     }
 
-    static String buildQueryShape(QueryBuilder queryBuilder) {
+    static String buildQueryShape(QueryBuilder queryBuilder, Boolean showFields) {
         if (queryBuilder == null) {
             return "";
         }
         QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
         queryBuilder.visit(shapeVisitor);
-        return shapeVisitor.prettyPrintTree("");
+        return shapeVisitor.prettyPrintTree("", showFields);
     }
 
     static String buildAggregationShape(AggregatorFactories.Builder aggregationsBuilder, Boolean showFields) {

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
@@ -18,17 +18,31 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.function.Function;
+
+import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.QUERY_FIELD_DATA_MAP;
+import static org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator.TWO_SPACE_INDENT;
 
 /**
  * Class to traverse the QueryBuilder tree and capture the query shape
  */
 public final class QueryShapeVisitor implements QueryBuilderVisitor {
     private final SetOnce<String> queryType = new SetOnce<>();
+    private final SetOnce<String> fieldData = new SetOnce<>();
     private final Map<BooleanClause.Occur, List<QueryShapeVisitor>> childVisitors = new EnumMap<>(BooleanClause.Occur.class);
 
     @Override
-    public void accept(QueryBuilder qb) {
-        queryType.set(qb.getName());
+    public void accept(QueryBuilder queryBuilder) {
+        queryType.set(queryBuilder.getName());
+
+        List<String> fieldDataList = new ArrayList<>();
+        List<Function<Object, String>> methods = QUERY_FIELD_DATA_MAP.get(queryBuilder.getClass());
+        if (methods != null) {
+            for (Function<Object, String> lambda : methods) {
+                fieldDataList.add(lambda.apply(queryBuilder));
+            }
+        }
+        fieldData.set(String.join(", ", fieldDataList));
     }
 
     @Override
@@ -84,12 +98,16 @@ public String toJson() {
      * @param indent indent size
      * @return Query builder tree as a pretty string
      */
-    public String prettyPrintTree(String indent) {
-        StringBuilder outputBuilder = new StringBuilder(indent).append(queryType.get()).append("\n");
+    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("\n");
         for (Map.Entry<BooleanClause.Occur, List<QueryShapeVisitor>> entry : childVisitors.entrySet()) {
-            outputBuilder.append(indent).append("  ").append(entry.getKey().name().toLowerCase(Locale.ROOT)).append(":\n");
+            outputBuilder.append(indent).append(TWO_SPACE_INDENT).append(entry.getKey().name().toLowerCase(Locale.ROOT)).append(":\n");
             for (QueryShapeVisitor child : entry.getValue()) {
-                outputBuilder.append(child.prettyPrintTree(indent + "    "));
+                outputBuilder.append(child.prettyPrintTree(indent + TWO_SPACE_INDENT.repeat(2), showFields));
             }
         }
         return outputBuilder.toString();

diff --git a/...ava/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeGeneratorTests.java b/...ava/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeGeneratorTests.java
@@ -11,6 +11,7 @@
 import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.MatchQueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.index.query.RangeQueryBuilder;
 import org.opensearch.index.query.RegexpQueryBuilder;
 import org.opensearch.index.query.TermQueryBuilder;
 import org.opensearch.plugin.insights.core.service.categorizer.QueryShapeGenerator;
@@ -24,15 +25,35 @@ public final class QueryShapeGeneratorTests extends OpenSearchTestCase {
     public void testQueryShape() {
         SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
         sourceBuilder.size(0);
-        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("field", "value2");
-        MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("tags", "php");
-        RegexpQueryBuilder regexpQueryBuilder = new RegexpQueryBuilder("field", "text");
-        sourceBuilder.query(new BoolQueryBuilder().must(termQueryBuilder).filter(matchQueryBuilder).should(regexpQueryBuilder));
+        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("field1", "value2");
+        MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("field2", "php");
+        RegexpQueryBuilder regexpQueryBuilder = new RegexpQueryBuilder("field3", "text");
+        RangeQueryBuilder rangeQueryBuilder = new RangeQueryBuilder("field4");
+        sourceBuilder.query(
+            new BoolQueryBuilder().must(termQueryBuilder).filter(matchQueryBuilder).should(regexpQueryBuilder).filter(rangeQueryBuilder)
+        );
 
-        String shape = QueryShapeGenerator.buildShape(sourceBuilder, true);
+        String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true);
+        String expectedShowFieldsTrue = "bool []\n"
+            + "  must:\n"
+            + "    term [field1]\n"
+            + "  filter:\n"
+            + "    match [field2]\n"
+            + "    range [field4]\n"
+            + "  should:\n"
+            + "    regexp [field3]\n";
+        assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);
 
-        String expected = "bool\n" + "  must:\n" + "    term\n" + "  filter:\n" + "    match\n" + "  should:\n" + "    regexp\n";
-        assertEquals(expected, shape);
+        String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false);
+        String expectedShowFieldsFalse = "bool\n"
+            + "  must:\n"
+            + "    term\n"
+            + "  filter:\n"
+            + "    match\n"
+            + "    range\n"
+            + "  should:\n"
+            + "    regexp\n";
+        assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse);
     }
 
     public void testAggregationShape() {
@@ -46,16 +67,25 @@ public void testAggregationShape() {
                 .subAggregation(new TermsAggregationBuilder("child-agg2").userValueTypeHint(ValueType.STRING).field("key.sub2"))
         );
 
-        String shape = QueryShapeGenerator.buildShape(sourceBuilder, true);
-
-        String expected = "aggregation:\n"
+        String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true);
+        String expectedShowFieldsTrue = "aggregation:\n"
             + "  terms [key]\n"
             + "    aggregation:\n"
             + "      terms [key.sub1]\n"
             + "      terms [key.sub2]\n"
             + "  terms [model]\n"
             + "  terms [type]\n";
-        assertEquals(expected, shape);
+        assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);
+
+        String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false);
+        String expectedShowFieldsFalse = "aggregation:\n"
+            + "  terms\n"
+            + "  terms\n"
+            + "  terms\n"
+            + "    aggregation:\n"
+            + "      terms\n"
+            + "      terms\n";
+        assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse);
     }
 
     public void testSortShape() {
@@ -65,9 +95,12 @@ public void testSortShape() {
         sourceBuilder.sort("price", SortOrder.ASC);
         sourceBuilder.sort("album", SortOrder.ASC);
 
-        String shape = QueryShapeGenerator.buildShape(sourceBuilder, true);
+        String shapeShowFieldsTrue = QueryShapeGenerator.buildShape(sourceBuilder, true);
+        String expectedShowFieldsTrue = "sort:\n" + "  asc [album]\n" + "  asc [price]\n" + "  desc [color]\n" + "  desc [vendor]\n";
+        assertEquals(expectedShowFieldsTrue, shapeShowFieldsTrue);
 
-        String expected = "sort:\n" + "  asc [album]\n" + "  asc [price]\n" + "  desc [color]\n" + "  desc [vendor]\n";
-        assertEquals(expected, shape);
+        String shapeShowFieldsFalse = QueryShapeGenerator.buildShape(sourceBuilder, false);
+        String expectedShowFieldsFalse = "sort:\n" + "  asc\n" + "  asc\n" + "  desc\n" + "  desc\n";
+        assertEquals(expectedShowFieldsFalse, shapeShowFieldsFalse);
     }
 }