Query shape for agg & sort

dzane17 · Jul 23, 2024 · cb4d238 · cb4d238
1 parent 29a3e2c
commit cb4d238
Show file tree

Hide file tree

Showing 2 changed files with 96 additions and 2 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizer.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizer.java
@@ -15,10 +15,13 @@
 import org.opensearch.index.query.QueryShapeVisitor;
 import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.sort.FieldSortBuilder;
 import org.opensearch.search.sort.SortBuilder;
 import org.opensearch.telemetry.metrics.MetricsRegistry;
 import org.opensearch.telemetry.metrics.tags.Tags;
 
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.ListIterator;
 
@@ -28,7 +31,8 @@
  */
 final class SearchQueryCategorizer {
 
-    private static final Logger log = LogManager.getLogger(SearchQueryCategorizer.class);
+    private static final Logger logger = LogManager.getLogger(SearchQueryCategorizer.class);
+    public static final String TWO_SPACE_INDENT = "  ";
 
     final SearchQueryCounters searchQueryCounters;
 
@@ -43,8 +47,12 @@ public void categorize(SearchSourceBuilder source) {
         QueryBuilder topLevelQueryBuilder = source.query();
         logQueryShape(topLevelQueryBuilder);
         incrementQueryTypeCounters(topLevelQueryBuilder);
+
         incrementQueryAggregationCounters(source.aggregations());
+        logAggregationsShape(source.aggregations());
+
         incrementQuerySortCounters(source.sorts());
+        logSortShape(source.sorts(), true);
     }
 
     private void incrementQuerySortCounters(List<SortBuilder<?>> sorts) {
@@ -79,7 +87,45 @@ private void logQueryShape(QueryBuilder topLevelQueryBuilder) {
         }
         QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
         topLevelQueryBuilder.visit(shapeVisitor);
-        log.trace("Query shape : {}", shapeVisitor.prettyPrintTree("  "));
+        logger.trace(shapeVisitor.prettyPrintTree(TWO_SPACE_INDENT));
+        System.out.println(shapeVisitor.prettyPrintTree(TWO_SPACE_INDENT));
     }
 
+    private void logAggregationsShape(AggregatorFactories.Builder aggregationsBuilder) {
+        if (aggregationsBuilder == null) {
+            return;
+        }
+        StringBuilder aggregationShape = aggregationsBuilder.visit(new StringBuilder(), 0, true);
+        logger.trace(aggregationShape.toString());
+        System.out.println(aggregationShape);
+    }
+
+    private void logSortShape(List<SortBuilder<?>> sortBuilderList, Boolean showFields) {
+        if (sortBuilderList.isEmpty()) {
+            return;
+        }
+        StringBuilder sortShape = new StringBuilder();
+        sortShape.append("sort:\n");
+
+        List<String> shapeStrings = new ArrayList<>();
+        for (SortBuilder sortBuilder : sortBuilderList) {
+            StringBuilder stringBuilder = new StringBuilder();
+            stringBuilder.append(TWO_SPACE_INDENT + sortBuilder.order() + ":");
+            if (showFields) {
+                if (sortBuilder instanceof FieldSortBuilder) {
+                    stringBuilder.append(" [" + ((FieldSortBuilder) sortBuilder).getFieldName() + "]");
+                } else {
+                    stringBuilder.append(" []");
+                }
+            }
+            shapeStrings.add(stringBuilder.toString());
+        }
+
+        Collections.sort(shapeStrings);
+        for (String line : shapeStrings) {
+            sortShape.append(line).append("\n");
+        }
+        logger.trace(sortShape.toString());
+        System.out.println(sortShape);
+    }
 }
diff --git a/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java b/server/src/main/java/org/opensearch/search/aggregations/AggregatorFactories.java
@@ -45,6 +45,8 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentLocation;
 import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.index.query.QueryBuilderVisitor;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.Rewriteable;
@@ -55,6 +57,7 @@
 import org.opensearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree;
 import org.opensearch.search.aggregations.support.AggregationPath;
 import org.opensearch.search.aggregations.support.AggregationPath.PathElement;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregationBuilder;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.profile.Profilers;
 import org.opensearch.search.profile.aggregation.ProfilingAggregator;
@@ -88,6 +91,7 @@
 @PublicApi(since = "1.0.0")
 public class AggregatorFactories {
     public static final Pattern VALID_AGG_NAME = Pattern.compile("[^\\[\\]>]+");
+    static final String TWO_SPACE_INDENT = "  ";
 
     /**
      * Parses the aggregation request recursively generating aggregator
@@ -372,6 +376,50 @@ public Builder(StreamInput in) throws IOException {
             }
         }
 
+        public StringBuilder visit(StringBuilder outputBuilder, int indentCount, Boolean showFields) {
+            String indent = TWO_SPACE_INDENT.repeat(indentCount);
+
+            //// Normal Aggregations ////
+            if (aggregationBuilders.isEmpty() == false) {
+                outputBuilder.append(indent + "aggregation:").append("\n");
+            }
+            List<String> aggShapeStrings = new ArrayList<>();
+            for (AggregationBuilder agg : aggregationBuilders) {
+                StringBuilder stringBuilder = new StringBuilder();
+                stringBuilder.append(indent + TWO_SPACE_INDENT + agg.getType());
+                if (showFields) {
+                    if (agg instanceof ValuesSourceAggregationBuilder) {
+                        stringBuilder.append(" [" + ((ValuesSourceAggregationBuilder) agg).field() + "]");
+                    } else {
+                        stringBuilder.append(" []");
+                    }
+                }
+                stringBuilder.append("\n");
+
+                if (agg.factoriesBuilder.names.isEmpty() == false) {
+                    // Recursive call on sub-aggregations
+                    stringBuilder = agg.factoriesBuilder.visit(stringBuilder, indentCount + 2, showFields);
+                }
+                aggShapeStrings.add(stringBuilder.toString());
+            }
+
+            // Sort aggregations
+            Collections.sort(aggShapeStrings);
+            for (String shapeString : aggShapeStrings) {
+                outputBuilder.append(shapeString);
+            }
+
+            //// Pipeline Aggregation (cannot have sub-aggregations) ////
+            if (pipelineAggregatorBuilders.isEmpty() == false) {
+                outputBuilder.append(indent + "pipeline aggregation:").append("\n");
+            }
+            for (PipelineAggregationBuilder agg : pipelineAggregatorBuilders) {
+                outputBuilder.append(indent + "  " + agg.name).append("\n");
+            }
+
+            return outputBuilder;
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeVInt(this.aggregationBuilders.size());