add ROUNDTO_PUSHDOWN_THRESHOLD in QueryParagmas and EsqlFlags to control the pushdown behavior at query and cluster level

Author: fang-xing-esql

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java

@@ -75,19 +75,11 @@ protected static List<Batch<PhysicalPlan>> rules(boolean optimizeForEsSource) {
         // execute the rules multiple times to improve the chances of things being pushed down
         @SuppressWarnings("unchecked")
         var pushdown = new Batch<PhysicalPlan>("Push to ES", esSourceRules.toArray(Rule[]::new));
-        List<Rule<?, PhysicalPlan>> substitutionRules = new ArrayList<>(1);
-        if (optimizeForEsSource) {
-            substitutionRules.add(new ReplaceRoundToWithQueryAndTags());
-        }
+
         // execute the SubstituteRoundToWithQueryAndTags rule once after all the other pushdown rules are applied, as this rule generate
         // multiple QueryBuilders according the number of RoundTo points, it should be applied after all the other eligible pushdowns are
         // done, and it should be executed only once.
-        @SuppressWarnings("unchecked")
-        var substituteRoundToWithQueryAndTags = new Batch<PhysicalPlan>(
-            "Substitute RoundTo with QueryAndTags",
-            Limiter.ONCE,
-            substitutionRules.toArray(Rule[]::new)
-        );
+        var substitutionRules = new Batch<>("Substitute RoundTo with QueryAndTags", Limiter.ONCE, new ReplaceRoundToWithQueryAndTags());
 
         // add the field extraction in just one pass
         // add it at the end after all the other rules have ran
@@ -99,6 +91,6 @@ protected static List<Batch<PhysicalPlan>> rules(boolean optimizeForEsSource) {
             new SpatialShapeBoundsExtraction(),
             new ParallelizeTimeSeriesSource()
         );
-        return List.of(pushdown, substituteRoundToWithQueryAndTags, fieldExtraction);
+        return optimizeForEsSource ? List.of(pushdown, substitutionRules, fieldExtraction) : List.of(pushdown, fieldExtraction);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTags.java

@@ -38,6 +38,7 @@
 
 import static org.elasticsearch.xpack.esql.core.type.DataTypeConverter.safeToLong;
 import static org.elasticsearch.xpack.esql.planner.TranslatorHandler.TRANSLATOR_HANDLER;
+import static org.elasticsearch.xpack.esql.plugin.QueryPragmas.ROUNDTO_PUSHDOWN_THRESHOLD;
 
 /**
  * {@code ReplaceRoundToWithQueryAndTags} builds a list of ranges and associated tags base on the rounding points defined in a
@@ -255,9 +256,6 @@ public class ReplaceRoundToWithQueryAndTags extends PhysicalOptimizerRules.Param
     EvalExec,
     LocalPhysicalOptimizerContext> {
 
-    // this is the maximum number of rounding points supported by this rule,
-    // it is the same as the maximum number of buckets used in Rounding.
-    public static final int MAX_NUM_POINTS = 127;
     private static final Logger logger = LogManager.getLogger(ReplaceRoundToWithQueryAndTags.class);
 
     @Override
@@ -277,12 +275,13 @@ protected PhysicalPlan rule(EvalExec evalExec, LocalPhysicalOptimizerContext ctx
             if (roundTos.size() == 1) {
                 RoundTo roundTo = roundTos.get(0);
                 int count = roundTo.points().size();
-                if (count > MAX_NUM_POINTS) {
+                int roundingPointsUpperLimit = roundingPointsThreshold(ctx);
+                if (count > roundingPointsUpperLimit) {
                     logger.debug(
                         "Skipping RoundTo push down for [{}], as it has [{}] points, which is more than [{}]",
                         roundTo.source(),
                         count,
-                        MAX_NUM_POINTS
+                        roundingPointsUpperLimit
                     );
                     return evalExec;
                 }
@@ -469,4 +468,21 @@ private static EsQueryExec.QueryBuilderAndTags buildCombinedQueryAndTags(
         QueryBuilder combinedQuery = Queries.combine(clause, mainQuery != null ? List.of(mainQuery, newQuery) : List.of(newQuery));
         return new EsQueryExec.QueryBuilderAndTags(combinedQuery, tags);
     }
+
+    /**
+     * Get the rounding points upper limit for {@code RoundTo} pushdown from query level pragmas or cluster level flags.
+     * If the query level pragmas is set to -1(default), the cluster level flags will be used.
+     * If the query level pragmas is set to greater than or equals to 0, the query level pragmas will be used.
+     */
+    private int roundingPointsThreshold(LocalPhysicalOptimizerContext ctx) {
+        int queryLevelRoundingPointsThreshold = ctx.configuration().pragmas().roundToPushDownThreshold();
+        int clusterLevelRoundingPointsThreshold = ctx.flags().roundToPushdownThreshold();
+        int roundingPointsThreshold;
+        if (queryLevelRoundingPointsThreshold == ROUNDTO_PUSHDOWN_THRESHOLD.getDefault(ctx.configuration().pragmas().getSettings())) {
+            roundingPointsThreshold = clusterLevelRoundingPointsThreshold;
+        } else {
+            roundingPointsThreshold = queryLevelRoundingPointsThreshold;
+        }
+        return roundingPointsThreshold;
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlFlags.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Setting;
+import org.elasticsearch.common.settings.Settings;
 
 public class EsqlFlags {
     public static final Setting<Boolean> ESQL_STRING_LIKE_ON_INDEX = Setting.boolSetting(
@@ -18,17 +19,65 @@ public class EsqlFlags {
         Setting.Property.Dynamic
     );
 
+    /**
+     * The maximum number of rounding points to push down to Lucene for the {@code roundTo} function at cluster level.
+     * {@code ReplaceRoundToWithQueryAndTags} checks this threshold before rewriting {@code RoundTo} to range queries.
+     *
+     * There is also a query level ROUNDTO_PUSHDOWN_THRESHOLD defined in {@code QueryPragmas}.
+     * The cluster level threshold defaults to 127, it is the same as the maximum number of buckets used in {@code Rounding}.
+     * The query level threshold defaults to -1, which means this query level setting is not set and cluster level upper limit will be used.
+     * If query level threshold is set to greater than or equals to 0, the query level threshold will be used, and it overrides the cluster
+     * level threshold.
+     *
+     * If the cluster level threshold is set to -1 or 0, no {@code RoundTo} pushdown will be performed, query level threshold is not set to
+     * -1 or 0.
+     */
+    public static final Setting<Integer> ESQL_ROUNDTO_PUSHDOWN_THRESHOLD = Setting.intSetting(
+        "esql.query.roundto_pushdown_threshold",
+        127,
+        -1,
+        Setting.Property.NodeScope,
+        Setting.Property.Dynamic
+    );
+
     private final boolean stringLikeOnIndex;
 
+    private final int roundToPushdownThreshold;
+
+    /**
+     * Constructor for tests.
+     */
     public EsqlFlags(boolean stringLikeOnIndex) {
         this.stringLikeOnIndex = stringLikeOnIndex;
+        this.roundToPushdownThreshold = ESQL_ROUNDTO_PUSHDOWN_THRESHOLD.getDefault(Settings.EMPTY);
+    }
+
+    /**
+     * Constructor for tests.
+     */
+    public EsqlFlags(int roundToPushdownThreshold) {
+        this.stringLikeOnIndex = ESQL_STRING_LIKE_ON_INDEX.getDefault(Settings.EMPTY);
+        this.roundToPushdownThreshold = roundToPushdownThreshold;
+    }
+
+    /**
+     * Constructor for tests.
+     */
+    public EsqlFlags(boolean stringLikeOnIndex, int roundToPushdownThreshold) {
+        this.stringLikeOnIndex = stringLikeOnIndex;
+        this.roundToPushdownThreshold = roundToPushdownThreshold;
     }
 
     public EsqlFlags(ClusterSettings settings) {
         this.stringLikeOnIndex = settings.get(ESQL_STRING_LIKE_ON_INDEX);
+        this.roundToPushdownThreshold = settings.get(ESQL_ROUNDTO_PUSHDOWN_THRESHOLD);
     }
 
     public boolean stringLikeOnIndex() {
         return stringLikeOnIndex;
     }
+
+    public int roundToPushdownThreshold() {
+        return roundToPushdownThreshold;
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java

@@ -260,7 +260,8 @@ public List<Setting<?>> getSettings() {
             PhysicalSettings.DEFAULT_DATA_PARTITIONING,
             PhysicalSettings.VALUES_LOADING_JUMBO_SIZE,
             STORED_FIELDS_SEQUENTIAL_PROPORTION,
-            EsqlFlags.ESQL_STRING_LIKE_ON_INDEX
+            EsqlFlags.ESQL_STRING_LIKE_ON_INDEX,
+            EsqlFlags.ESQL_ROUNDTO_PUSHDOWN_THRESHOLD
         );
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/QueryPragmas.java

@@ -80,6 +80,20 @@ public final class QueryPragmas implements Writeable {
         MappedFieldType.FieldExtractPreference.NONE
     );
 
+    /**
+     * The maximum number of rounding points to push down to Lucene for the {@code roundTo} function at query level.
+     * {@code ReplaceRoundToWithQueryAndTags} checks this threshold before rewriting {@code RoundTo} to range queries.
+     *
+     * There is also a cluster level ESQL_ROUNDTO_PUSHDOWN_THRESHOLD defined in {@code EsqlFlags}.
+     * The query level threshold defaults to -1, which means this query level setting is not set and cluster level upper limit will be used.
+     * The cluster level threshold defaults to 127, it is the same as the maximum number of buckets used in {@code Rounding}.
+     * If query level threshold is set to greater than or equals to 0, the query level threshold will be used, and it overrides the cluster
+     * level threshold.
+     *
+     * If the query level threshold is set to 0, no {@code RoundTo} pushdown will be performed.
+     */
+    public static final Setting<Integer> ROUNDTO_PUSHDOWN_THRESHOLD = Setting.intSetting("roundto_pushdown_threshold", -1, -1);
+
     public static final QueryPragmas EMPTY = new QueryPragmas(Settings.EMPTY);
 
     private final Settings settings;
@@ -197,6 +211,10 @@ public MappedFieldType.FieldExtractPreference fieldExtractPreference() {
         return FIELD_EXTRACT_PREFERENCE.get(settings);
     }
 
+    public int roundToPushDownThreshold() {
+        return ROUNDTO_PUSHDOWN_THRESHOLD.get(settings);
+    }
+
     public boolean isEmpty() {
         return settings.isEmpty();
     }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/TestPlannerOptimizer.java

@@ -55,7 +55,15 @@ public PhysicalPlan plan(String query, SearchStats stats, Analyzer analyzer) {
         return physical;
     }
 
+    public PhysicalPlan plan(String query, SearchStats stats, EsqlFlags esqlFlags) {
+        return optimizedPlan(physicalPlan(query, analyzer), stats, esqlFlags);
+    }
+
     private PhysicalPlan optimizedPlan(PhysicalPlan plan, SearchStats searchStats) {
+        return optimizedPlan(plan, searchStats, new EsqlFlags(true));
+    }
+
+    private PhysicalPlan optimizedPlan(PhysicalPlan plan, SearchStats searchStats, EsqlFlags esqlFlags) {
         // System.out.println("* Physical Before\n" + plan);
         var physicalPlan = EstimatesRowSize.estimateRowSize(0, physicalPlanOptimizer.optimize(plan));
         // System.out.println("* Physical After\n" + physicalPlan);
@@ -67,7 +75,7 @@ private PhysicalPlan optimizedPlan(PhysicalPlan plan, SearchStats searchStats) {
             new LocalLogicalOptimizerContext(config, FoldContext.small(), searchStats)
         );
         var physicalTestOptimizer = new TestLocalPhysicalPlanOptimizer(
-            new LocalPhysicalOptimizerContext(new EsqlFlags(true), config, FoldContext.small(), searchStats),
+            new LocalPhysicalOptimizerContext(esqlFlags, config, FoldContext.small(), searchStats),
             true
         );
         var l = PlannerUtils.localPlan(physicalPlan, logicalTestOptimizer, physicalTestOptimizer);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTagsTests.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.optimizer.rules.physical.local;
 
 import org.elasticsearch.common.logging.LoggerMessageFormat;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.index.query.BoolQueryBuilder;
 import org.elasticsearch.index.query.MatchQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -27,6 +28,7 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo;
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizerTests;
+import org.elasticsearch.xpack.esql.optimizer.TestPlannerOptimizer;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsQueryExec;
 import org.elasticsearch.xpack.esql.plan.physical.EvalExec;
@@ -38,12 +40,15 @@
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.plan.physical.ProjectExec;
 import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
+import org.elasticsearch.xpack.esql.plugin.EsqlFlags;
+import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 import org.elasticsearch.xpack.esql.session.Configuration;
 import org.elasticsearch.xpack.esql.stats.SearchStats;
 
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.stream.Collectors;
 
@@ -55,6 +60,7 @@
 import static org.elasticsearch.index.query.QueryBuilders.rangeQuery;
 import static org.elasticsearch.index.query.QueryBuilders.termQuery;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.configuration;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.DEFAULT_DATE_NANOS_FORMATTER;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.DEFAULT_DATE_TIME_FORMATTER;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateNanosToLong;
@@ -391,9 +397,11 @@ public void testDateTruncBucketNotTransformToQueryAndTagsWithFork() {
         }
     }
 
-    // If the number of rounding points is 127 or less, the query is rewritten to use QueryAndTags.
-    // If the number of rounding points is 128 or more, the query is not rewritten.
-    public void testRoundToTransformToQueryAndTagsUpperLimit() {
+    /**
+     * If the number of rounding points is 127 or less, the query is rewritten to QueryAndTags.
+     * If the number of rounding points is 128 or more, the query is not rewritten.
+     */
+    public void testRoundToTransformToQueryAndTagsWithDefaultUpperLimit() {
         for (int numOfPoints : List.of(127, 128)) {
             StringBuilder points = new StringBuilder();
             for (int i = 0; i < numOfPoints; i++) {
@@ -445,6 +453,83 @@ public void testRoundToTransformToQueryAndTagsUpperLimit() {
         }
     }
 
+    /**
+     * Query level threshold(if greater than -1) set in QueryPragmas overrides the cluster level threshold set in EsqlFlags.
+     */
+    public void testRoundToTransformToQueryAndTagsWithCustomizedUpperLimit() {
+        for (int clusterLevelThreshold : List.of(-1, 0, 60, 126, 128, 256)) {
+            for (int queryLevelThreshold : List.of(-1, 0, 60, 126, 128, 256)) {
+                StringBuilder points = new StringBuilder(); // there are 127 rounding points
+                for (int i = 0; i < 127; i++) {
+                    if (i > 0) {
+                        points.append(", ");
+                    }
+                    points.append(i);
+                }
+                String query = LoggerMessageFormat.format(null, """
+                    from test
+                    | stats count(*) by x = round_to(integer, {})
+                    """, points.toString());
+
+                TestPlannerOptimizer plannerOptimizerWithPragmas = new TestPlannerOptimizer(
+                    configuration(
+                        new QueryPragmas(
+                            Settings.builder()
+                                .put(QueryPragmas.ROUNDTO_PUSHDOWN_THRESHOLD.getKey().toLowerCase(Locale.ROOT), queryLevelThreshold)
+                                .build()
+                        ),
+                        query
+                    ),
+                    makeAnalyzer("mapping-all-types.json")
+                );
+                EsqlFlags esqlFlags = new EsqlFlags(clusterLevelThreshold);
+                assertEquals(clusterLevelThreshold, esqlFlags.roundToPushdownThreshold());
+                assertTrue(esqlFlags.stringLikeOnIndex());
+                PhysicalPlan plan = plannerOptimizerWithPragmas.plan(query, searchStats, esqlFlags);
+                boolean pushdown = false;
+                if (queryLevelThreshold > -1) {
+                    pushdown = queryLevelThreshold >= 127;
+                } else {
+                    pushdown = clusterLevelThreshold >= 127;
+                }
+
+                LimitExec limit = as(plan, LimitExec.class);
+                AggregateExec agg = as(limit.child(), AggregateExec.class);
+                assertThat(agg.getMode(), is(FINAL));
+                List<? extends Expression> groupings = agg.groupings();
+                NamedExpression grouping = as(groupings.get(0), NamedExpression.class);
+                assertEquals("x", grouping.name());
+                assertEquals(DataType.INTEGER, grouping.dataType());
+                assertEquals(List.of("count(*)", "x"), Expressions.names(agg.aggregates()));
+                ExchangeExec exchange = as(agg.child(), ExchangeExec.class);
+                assertThat(exchange.inBetweenAggs(), is(true));
+                agg = as(exchange.child(), AggregateExec.class);
+                EvalExec evalExec = as(agg.child(), EvalExec.class);
+                List<Alias> aliases = evalExec.fields();
+                assertEquals(1, aliases.size());
+                if (pushdown) {
+                    FieldAttribute roundToTag = as(aliases.get(0).child(), FieldAttribute.class);
+                    assertTrue(roundToTag.name().startsWith("$$integer$round_to$"));
+                    EsQueryExec esQueryExec = as(evalExec.child(), EsQueryExec.class);
+                    List<EsQueryExec.QueryBuilderAndTags> queryBuilderAndTags = esQueryExec.queryBuilderAndTags();
+                    assertEquals(128, queryBuilderAndTags.size()); // 127 + nullBucket
+                    assertThrows(UnsupportedOperationException.class, esQueryExec::query);
+                } else { // query rewrite does not happen
+                    RoundTo roundTo = as(aliases.get(0).child(), RoundTo.class);
+                    assertEquals(127, roundTo.points().size());
+                    FieldExtractExec fieldExtractExec = as(evalExec.child(), FieldExtractExec.class);
+                    EsQueryExec esQueryExec = as(fieldExtractExec.child(), EsQueryExec.class);
+                    List<EsQueryExec.QueryBuilderAndTags> queryBuilderAndTags = esQueryExec.queryBuilderAndTags();
+                    assertEquals(1, queryBuilderAndTags.size());
+                    EsQueryExec.QueryBuilderAndTags queryBuilder = queryBuilderAndTags.get(0);
+                    assertNull(queryBuilder.query());
+                    assertTrue(queryBuilder.tags().isEmpty());
+                    assertNull(esQueryExec.query());
+                }
+            }
+        }
+    }
+
     private static void verifyQueryAndTags(List<EsQueryExec.QueryBuilderAndTags> expected, List<EsQueryExec.QueryBuilderAndTags> actual) {
         assertEquals(expected.size(), actual.size());
         for (int i = 0; i < expected.size(); i++) {