upper limit on the number of rounding points when converting roundTo to range queries

Author: fang-xing-esql

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

@@ -8,6 +8,8 @@
 package org.elasticsearch.xpack.esql.optimizer.rules.physical.local;
 
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.logging.Logger;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
@@ -253,6 +255,11 @@ 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
     protected PhysicalPlan rule(EvalExec evalExec, LocalPhysicalOptimizerContext ctx) {
         PhysicalPlan plan = evalExec;
@@ -268,7 +275,18 @@ protected PhysicalPlan rule(EvalExec evalExec, LocalPhysicalOptimizerContext ctx
                 .toList();
             // It is not clear how to push down multiple RoundTos, dealing with multiple RoundTos is out of the scope of this PR.
             if (roundTos.size() == 1) {
-                plan = planRoundTo(roundTos.get(0), evalExec, queryExec, ctx);
+                RoundTo roundTo = roundTos.get(0);
+                int count = roundTo.points().size();
+                if (count > MAX_NUM_POINTS) {
+                    logger.debug(
+                        "Skipping RoundTo push down for [{}], as it has [{}] points, which is more than [{}]",
+                        roundTo.source(),
+                        count,
+                        MAX_NUM_POINTS
+                    );
+                    return evalExec;
+                }
+                plan = planRoundTo(roundTo, evalExec, queryExec, ctx);
             }
         }
         return plan;

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

@@ -391,6 +391,60 @@ 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() {
+        for (int numOfPoints : List.of(127, 128)) {
+            StringBuilder points = new StringBuilder();
+            for (int i = 0; i < numOfPoints; 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());
+
+            PhysicalPlan plan = plannerOptimizer.plan(query, searchStats, makeAnalyzer("mapping-all-types.json"));
+
+            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 (numOfPoints == 127) {
+                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 { // numOfPoints == 128, query rewrite does not happen
+                RoundTo roundTo = as(aliases.get(0).child(), RoundTo.class);
+                assertEquals(128, 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++) {