[ES|QL] Rewrite RoundTo to QueryAndTags (elastic#132512)

* Rewrite RoundTo to QueryAndTags

Author: fang-xing-esql

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/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundToDouble.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.core.type.DataTypeConverter;
 import org.elasticsearch.xpack.esql.expression.Foldables;
 import org.elasticsearch.xpack.esql.expression.function.Example;
 import org.elasticsearch.xpack.esql.expression.function.FunctionAppliesTo;
@@ -30,6 +31,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;
@@ -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,27 @@ interface Build {
         Map.entry(LONG, RoundToLong.BUILD),
         Map.entry(DOUBLE, RoundToDouble.BUILD)
     );
+
+    public static List<Object> sortedRoundingPoints(List<Object> points, DataType dataType) {
+        List<Number> pointsTobeSorted = points.stream().filter(Objects::nonNull).map(p -> (Number) p).toList();
+
+        return switch (dataType) {
+            case INTEGER -> pointsTobeSorted.stream()
+                .mapToInt(Number::intValue)
+                .sorted()
+                .boxed()
+                .collect(java.util.stream.Collectors.toList());
+            case DOUBLE -> pointsTobeSorted.stream()
+                .mapToDouble(Number::doubleValue)
+                .sorted()
+                .boxed()
+                .collect(java.util.stream.Collectors.toList());
+            case LONG, DATETIME, DATE_NANOS -> pointsTobeSorted.stream()
+                .mapToLong(DataTypeConverter::safeToLong)
+                .sorted()
+                .boxed()
+                .collect(java.util.stream.Collectors.toList());
+            default -> throw new IllegalArgumentException("Unsupported data type: " + dataType);
+        };
+    }
 }

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

@@ -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/generated-src/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundToLong.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,12 @@ 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));
+
+        // 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.
+        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
         var fieldExtraction = new Batch<>(
@@ -84,6 +91,6 @@ protected static List<Batch<PhysicalPlan>> rules(boolean optimizeForEsSource) {
             new SpatialShapeBoundsExtraction(),
             new ParallelizeTimeSeriesSource()
         );
-        return List.of(pushdown, fieldExtraction);
+        return optimizeForEsSource ? List.of(pushdown, substitutionRules, fieldExtraction) : List.of(pushdown, fieldExtraction);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/math/RoundTo.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);