Address comments, add agg and sort counters, add feature flag, refact…

…oring Signed-off-by: Siddhant Deshmukh <[email protected]>
deshsidd · Oct 11, 2023 · 5b7c17d · 5b7c17d
1 parent fb6438e
commit 5b7c17d
Show file tree

Hide file tree

Showing 22 changed files with 177 additions and 107 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizingVisitor.java
@@ -0,0 +1,71 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.action.search;
+
+import org.apache.lucene.search.BooleanClause;
+import org.opensearch.index.query.BoolQueryBuilder;
+import org.opensearch.index.query.MatchPhraseQueryBuilder;
+import org.opensearch.index.query.MatchQueryBuilder;
+import org.opensearch.index.query.MultiMatchQueryBuilder;
+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.TermQueryBuilder;
+import org.opensearch.index.query.WildcardQueryBuilder;
+import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
+import org.opensearch.telemetry.metrics.tags.Tags;
+
+/**
+ * Class to visit the querybuilder tree and also track the level information.
+ * Increments the counters related to Search Query type.
+ */
+public class SearchQueryCategorizingVisitor implements QueryBuilderVisitor {
+    private final int level;
+    private final SearchQueryCounters searchQueryCounters;
+
+    public SearchQueryCategorizingVisitor(SearchQueryCounters searchQueryCounters) {
+        this(searchQueryCounters, 0);
+    }
+
+    private SearchQueryCategorizingVisitor(SearchQueryCounters counters, int level) {
+        this.searchQueryCounters = counters;
+        this.level = level;
+    }
+    public void accept(QueryBuilder qb) {
+        if (qb instanceof BoolQueryBuilder) {
+            searchQueryCounters.boolCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof FunctionScoreQueryBuilder) {
+            searchQueryCounters.functionScoreCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof MatchQueryBuilder) {
+            searchQueryCounters.matchCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof MatchPhraseQueryBuilder) {
+            searchQueryCounters.matchPhrasePrefixCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof MultiMatchQueryBuilder) {
+            searchQueryCounters.multiMatchCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof QueryStringQueryBuilder) {
+            searchQueryCounters.queryStringQueryCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof RangeQueryBuilder) {
+            searchQueryCounters.rangeCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof RegexpQueryBuilder) {
+            searchQueryCounters.regexCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof TermQueryBuilder) {
+            searchQueryCounters.termCounter.add(1, Tags.create().addTag("level", level));
+        } else if (qb instanceof WildcardQueryBuilder) {
+            searchQueryCounters.wildcardCounter.add(1, Tags.create().addTag("level", level));
+        } else {
+            searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag("level", level));
+        }
+    }
+
+    public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
+        return new SearchQueryCategorizingVisitor(searchQueryCounters, level + 1);
+    }
+}
diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java
@@ -10,24 +10,17 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
-import org.apache.lucene.search.BooleanClause;
-import org.opensearch.index.query.BoolQueryBuilder;
-import org.opensearch.index.query.MatchPhraseQueryBuilder;
-import org.opensearch.index.query.MatchQueryBuilder;
-import org.opensearch.index.query.MultiMatchQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilderVisitor;
 import org.opensearch.index.query.QueryShapeVisitor;
-import org.opensearch.index.query.QueryStringQueryBuilder;
-import org.opensearch.index.query.RangeQueryBuilder;
-import org.opensearch.index.query.RegexpQueryBuilder;
-import org.opensearch.index.query.TermQueryBuilder;
-import org.opensearch.index.query.WildcardQueryBuilder;
-import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
+import org.opensearch.search.aggregations.AggregatorFactories;
 import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.sort.SortBuilder;
 import org.opensearch.telemetry.metrics.MetricsRegistry;
 import org.opensearch.telemetry.metrics.tags.Tags;
 
+import java.util.List;
+
 /**
  * Class to categorize the search queries based on the type and increment the relevant counters.
  * Class also logs the query shape.
@@ -46,51 +39,31 @@ public void categorize(SearchSourceBuilder source) {
         QueryBuilder topLevelQueryBuilder = source.query();
 
         logQueryShape(topLevelQueryBuilder);
+        incrementQueryTypeCounters(topLevelQueryBuilder);
+        incrementQueryAggregationCounters(source.aggregations());
+        incrementQuerySortCounters(source.sorts());
+    }
 
-        incrementQueryCounters(topLevelQueryBuilder);
+    private void incrementQuerySortCounters(List<SortBuilder<?>> sorts) {
+        if (sorts.size() > 0) {
+            searchQueryCounters.sortCounter.add(1);
+        }
     }
 
-    private static void incrementQueryCounters(QueryBuilder topLevelQueryBuilder) {
-        // Increment the query counters using Metric Framework
-        QueryBuilderVisitor queryBuilderVisitor = new QueryBuilderVisitor() {
-            @Override
-            public void accept(QueryBuilder qb, int level) {
-                if (qb instanceof BoolQueryBuilder) {
-                    searchQueryCounters.boolCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof FunctionScoreQueryBuilder) {
-                    searchQueryCounters.functionScoreCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof MatchQueryBuilder) {
-                    searchQueryCounters.matchCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof MatchPhraseQueryBuilder) {
-                    searchQueryCounters.matchPhrasePrefixCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof MultiMatchQueryBuilder) {
-                    searchQueryCounters.multiMatchCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof QueryStringQueryBuilder) {
-                    searchQueryCounters.queryStringQueryCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof RangeQueryBuilder) {
-                    searchQueryCounters.rangeCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof RegexpQueryBuilder) {
-                    searchQueryCounters.regexCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof TermQueryBuilder) {
-                    searchQueryCounters.termCounter.add(1, Tags.create().addTag("level", level));
-                } else if (qb instanceof WildcardQueryBuilder) {
-                    searchQueryCounters.wildcardCounter.add(1, Tags.create().addTag("level", level));
-                } else {
-                    searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag("level", level));
-                }
-            }
+    private void incrementQueryAggregationCounters(AggregatorFactories.Builder aggregations) {
+        if (aggregations != null) {
+            searchQueryCounters.aggCounter.add(1);
+        }
+    }
 
-            @Override
-            public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
-                return this;
-            }
-        };
-        topLevelQueryBuilder.visit(queryBuilderVisitor, 0);
+    private static void incrementQueryTypeCounters(QueryBuilder topLevelQueryBuilder) {
+        QueryBuilderVisitor searhQueryVisitor = new SearchQueryCategorizingVisitor(searchQueryCounters);
+        topLevelQueryBuilder.visit(searhQueryVisitor);
     }
 
     private static void logQueryShape(QueryBuilder topLevelQueryBuilder) {
         QueryShapeVisitor shapeVisitor = new QueryShapeVisitor();
-        topLevelQueryBuilder.visit(shapeVisitor, 0);
+        topLevelQueryBuilder.visit(shapeVisitor);
         String queryShapeJson = shapeVisitor.prettyPrintTree("  ");
         log.debug("Query shape : " + queryShapeJson);
     }

diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java
@@ -28,6 +28,8 @@ public class SearchQueryCounters {
     public final Counter queryStringQueryCounter;
     public final Counter rangeCounter;
     public final Counter regexCounter;
+
+    public final Counter sortCounter;
     public final Counter termCounter;
     public final Counter totalCounter;
     public final Counter wildcardCounter;
@@ -36,7 +38,7 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) {
         this.metricsRegistry = metricsRegistry;
         this.aggCounter = metricsRegistry.createCounter(
             "aggSearchQueryCounter",
-            "Counter for the number of top level and nested agg search queries",
+            "Counter for the number of top level agg search queries",
             "0"
         );
         this.boolCounter = metricsRegistry.createCounter(
@@ -84,6 +86,11 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) {
             "Counter for the number of top level and nested regex search queries",
             "0"
         );
+        this.sortCounter = metricsRegistry.createCounter(
+            "sortSearchQueryCounter",
+            "Counter for the number of top level sort search queries",
+            "0"
+        );
         this.termCounter = metricsRegistry.createCounter(
             "termSearchQueryCounter",
             "Counter for the number of top level and nested term search queries",

diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -60,6 +60,7 @@
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
 import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.util.concurrent.AtomicArray;
 import org.opensearch.common.util.concurrent.CountDown;
 import org.opensearch.core.action.ActionListener;
@@ -173,6 +174,8 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,
 
     private final MetricsRegistry metricsRegistry;
 
+    private final SearchQueryCategorizor searchQueryCategorizor;
+
     @Inject
     public TransportSearchAction(
         NodeClient client,
@@ -207,6 +210,11 @@ public TransportSearchAction(
         clusterService.getClusterSettings().addSettingsUpdateConsumer(SEARCH_REQUEST_STATS_ENABLED, this::setIsRequestStatsEnabled);
         this.searchRequestStats = searchRequestStats;
         this.metricsRegistry = metricsRegistry;
+        if (FeatureFlags.isEnabled(FeatureFlags.QUERY_CATEOGORIZATION)) {
+            this.searchQueryCategorizor = new SearchQueryCategorizor(metricsRegistry);
+        } else {
+            this.searchQueryCategorizor = null;
+        }
     }
 
     private void setIsRequestStatsEnabled(boolean isRequestStatsEnabled) {
@@ -437,8 +445,9 @@ private void executeRequest(
             return;
         }
 
-        SearchQueryCategorizor searchQueryCategorizor = new SearchQueryCategorizor(metricsRegistry);
-        searchQueryCategorizor.categorize(searchRequest.source());
+        if (FeatureFlags.isEnabled(FeatureFlags.QUERY_CATEOGORIZATION)) {
+            searchQueryCategorizor.categorize(searchRequest.source());
+        }
 
         ActionListener<SearchSourceBuilder> rewriteListener = ActionListener.wrap(source -> {
             if (source != searchRequest.source()) {

diff --git a/server/src/main/java/org/opensearch/common/util/FeatureFlags.java b/server/src/main/java/org/opensearch/common/util/FeatureFlags.java
@@ -60,6 +60,11 @@ public class FeatureFlags {
      */
     public static final String DATETIME_FORMATTER_CACHING = "opensearch.experimental.optimization.datetime_formatter_caching.enabled";
 
+    /**
+     * Gates the query categorization by type feature.
+     */
+    public static final String QUERY_CATEOGORIZATION = "opensearch.experimental.feature.query_categorization.enabled";
+
     /**
      * Should store the settings from opensearch.yml.
      */
@@ -122,4 +127,10 @@ public static boolean isEnabled(Setting<Boolean> featureFlag) {
         true,
         Property.NodeScope
     );
+
+    public static final Setting<Boolean> QUERY_CATEGORIZATION_SETTING = Setting.boolSetting(
+        QUERY_CATEOGORIZATION,
+        false,
+        Property.NodeScope
+    );
 }
diff --git a/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoolQueryBuilder.java
@@ -429,30 +429,30 @@ private static boolean rewriteClauses(
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor, int level) {
-        visitor.accept(this, level);
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
         if (mustClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST);
             for (QueryBuilder mustClause : mustClauses) {
-                mustClause.visit(subVisitor, level + 1);
+                mustClause.visit(subVisitor);
             }
         }
         if (shouldClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.SHOULD);
             for (QueryBuilder shouldClause : shouldClauses) {
-                shouldClause.visit(subVisitor, level + 1);
+                shouldClause.visit(subVisitor);
             }
         }
         if (mustNotClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST_NOT);
             for (QueryBuilder mustNotClause : mustNotClauses) {
-                mustNotClause.visit(subVisitor, level + 1);
+                mustNotClause.visit(subVisitor);
             }
         }
         if (filterClauses.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.FILTER);
             for (QueryBuilder filterClause : filterClauses) {
-                filterClause.visit(subVisitor, level + 1);
+                filterClause.visit(subVisitor);
             }
         }
 

diff --git a/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/BoostingQueryBuilder.java
@@ -255,13 +255,13 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor, int level) {
-        visitor.accept(this, level);
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
         if (positiveQuery != null) {
-            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(positiveQuery, level + 1);
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(positiveQuery);
         }
         if (negativeQuery != null) {
-            visitor.getChildVisitor(BooleanClause.Occur.SHOULD).accept(negativeQuery, level + 1);
+            visitor.getChildVisitor(BooleanClause.Occur.SHOULD).accept(negativeQuery);
         }
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/ConstantScoreQueryBuilder.java
@@ -186,9 +186,9 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor, int level) {
-        visitor.accept(this, level);
-        visitor.getChildVisitor(BooleanClause.Occur.FILTER).accept(filterBuilder, level + 1);
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.FILTER).accept(filterBuilder);
     }
 
 }
diff --git a/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/DisMaxQueryBuilder.java
@@ -249,12 +249,12 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor, int level) {
-        visitor.accept(this, level);
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
         if (queries.isEmpty() == false) {
             QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.SHOULD);
             for (QueryBuilder subQb : queries) {
-                subVisitor.accept(subQb, level + 1);
+                subVisitor.accept(subQb);
             }
         }
     }

diff --git a/server/src/main/java/org/opensearch/index/query/FieldMaskingSpanQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/FieldMaskingSpanQueryBuilder.java
@@ -210,8 +210,8 @@ public String getWriteableName() {
     }
 
     @Override
-    public void visit(QueryBuilderVisitor visitor, int level) {
-        visitor.accept(this, level);
-        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(queryBuilder, level + 1);
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(queryBuilder);
     }
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilder.java b/server/src/main/java/org/opensearch/index/query/QueryBuilder.java
@@ -99,10 +99,9 @@ default QueryBuilder rewrite(QueryRewriteContext queryShardContext) throws IOExc
     /**
      * Recurse through the QueryBuilder tree, visiting any child QueryBuilder.
      * @param visitor a query builder visitor to be called by each query builder in the tree.
-     * @param level level of the current query builder based on the QueryBuilder tree. 0 is top/root level.
      */
-    default void visit(QueryBuilderVisitor visitor, int level) {
-        visitor.accept(this, level);
+    default void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
     };
 
 }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java b/server/src/main/java/org/opensearch/index/query/QueryBuilderVisitor.java
@@ -18,9 +18,8 @@ public interface QueryBuilderVisitor {
     /**
      * Accept method is called when the visitor accepts the queryBuilder object to be traversed in the query tree.
      * @param qb is a queryBuilder object which is accepeted by the visitor.
-     * @param level level of the current query builder based on the QueryBuilder tree. 0 is top/root level.
      */
-    void accept(QueryBuilder qb, int level);
+    void accept(QueryBuilder qb);
 
     /**
      * Fetches the child sub visitor from the main QueryBuilderVisitor Object.
@@ -36,7 +35,7 @@ public interface QueryBuilderVisitor {
      */
     QueryBuilderVisitor NO_OP_VISITOR = new QueryBuilderVisitor() {
         @Override
-        public void accept(QueryBuilder qb, int level) {
+        public void accept(QueryBuilder qb) {
             // Do nothing
         }