[ES|QL] Rewrite RoundTo to QueryAndTags

docs/changelog/132512.yaml

@@ -0,0 +1,5 @@
+pr: 132512
+summary: Rewrite `RoundTo` to `QueryAndTags`
+area: ES|QL
+type: enhancement
+issues: []

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundTo.java

@@ -30,6 +30,7 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isFoldable;
@@ -39,6 +40,7 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
 import static org.elasticsearch.xpack.esql.core.type.DataType.LONG;
 import static org.elasticsearch.xpack.esql.core.type.DataType.UNSIGNED_LONG;
+import static org.elasticsearch.xpack.esql.core.type.DataTypeConverter.safeToLong;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.commonType;
 
 /**
@@ -181,7 +183,8 @@ public ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
         ExpressionEvaluator.Factory field = toEvaluator.apply(field());
         field = Cast.cast(source(), field().dataType(), dataType, field);
         List<Object> points = Iterators.toList(Iterators.map(points().iterator(), p -> Foldables.valueOf(toEvaluator.foldCtx(), p)));
-        return build.build(source(), field, points);
+        List<Object> sortedPoints = sortedRoundingPoints(points, dataType); // provide sorted points to the evaluator
+        return build.build(source(), field, sortedPoints);
     }
 
     interface Build {
@@ -195,4 +198,19 @@ interface Build {
         Map.entry(LONG, RoundToLong.BUILD),
         Map.entry(DOUBLE, RoundToDouble.BUILD)
     );
+
+    public static List<Object> sortedRoundingPoints(List<Object> points, DataType dataType) {
+        return points.stream().filter(Objects::nonNull).map(p -> switch (dataType) { // the types are in sync with SIGNATURES
+            case INTEGER -> ((Number) p).intValue();
+            case DOUBLE -> ((Number) p).doubleValue();
+            case LONG, DATETIME, DATE_NANOS -> safeToLong((Number) p);
+            default -> throw new IllegalArgumentException("Unsupported data type: " + dataType);
+        }).sorted((a, b) -> {
+            if (a instanceof Double || b instanceof Double) {
+                return Double.compare(a.doubleValue(), b.doubleValue());
+            } else {
+                return Long.compare(a.longValue(), b.longValue());
+            }
+        }).collect(java.util.stream.Collectors.toList());
+    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/X-RoundTo.java.st

@@ -25,7 +25,6 @@ import java.util.Arrays;
 class RoundTo$Type$ {
     static final RoundTo.Build BUILD = (source, field, points) -> {
         $type$[] f = points.stream().mapTo$Type$(p -> ((Number) p).$type$Value()).toArray();
-        Arrays.sort(f);
         return switch (f.length) {
             // TODO should be a consistent way to do the 0 version - is CASE(MV_COUNT(f) == 1, f[0])
             case 1 -> new RoundTo$Type$1Evaluator.Factory(source, field, f[0]);

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

@@ -18,6 +18,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushSampleToSource;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushStatsToSource;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushTopNToSource;
+import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.ReplaceRoundToWithQueryAndTags;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.ReplaceSourceAttributes;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.SpatialDocValuesExtraction;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.SpatialShapeBoundsExtraction;
@@ -60,7 +61,7 @@ protected List<Batch<PhysicalPlan>> batches() {
     }
 
     protected static List<Batch<PhysicalPlan>> rules(boolean optimizeForEsSource) {
-        List<Rule<?, PhysicalPlan>> esSourceRules = new ArrayList<>(6);
+        List<Rule<?, PhysicalPlan>> esSourceRules = new ArrayList<>(7);
         esSourceRules.add(new ReplaceSourceAttributes());
         if (optimizeForEsSource) {
             esSourceRules.add(new PushTopNToSource());
@@ -74,6 +75,20 @@ 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)
+        );
+
         // add the field extraction in just one pass
         // add it at the end after all the other rules have ran
         var fieldExtraction = new Batch<>(
@@ -84,6 +99,6 @@ protected static List<Batch<PhysicalPlan>> rules(boolean optimizeForEsSource) {
             new SpatialShapeBoundsExtraction(),
             new ParallelizeTimeSeriesSource()
         );
-        return List.of(pushdown, fieldExtraction);
+        return List.of(pushdown, substituteRoundToWithQueryAndTags, fieldExtraction);
     }
 }

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

@@ -126,10 +126,10 @@ private static PhysicalPlan rewrite(
                 queryExec.indexMode(),
                 queryExec.indexNameWithModes(),
                 queryExec.output(),
-                query,
                 queryExec.limit(),
                 queryExec.sorts(),
-                queryExec.estimatedRowSize()
+                queryExec.estimatedRowSize(),
+                List.of(new EsQueryExec.QueryBuilderAndTags(query, List.of()))
             );
             // If the eval contains other aliases, not just field attributes, we need to keep them in the plan
             PhysicalPlan plan = evalFields.isEmpty() ? queryExec : new EvalExec(filterExec.source(), queryExec, evalFields);