refactor according to review comments

Author: fang-xing-esql

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

@@ -0,0 +1,28 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.expression;
+
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.stats.SearchStats;
+
+/**
+ * Interface signaling to the local logical plan optimizer that the declaring expression
+ * has to be replaced by a different form.
+ * Implement this on {@code Function}s when:
+ * <ul>
+ *     <li>The expression can be rewritten to another expression on data node, with the statistics available in SearchStats.
+ *     Like {@code DateTrunc} and {@code Bucket} could be rewritten to {@code RoundTo} with the min/max values on the date field.
+ *     </li>
+ * </ul>
+ */
+public interface LocalSurrogateExpression {
+    /**
+     * Returns the expression to be replaced by or {@code null} if this cannot be replaced.
+     */
+    Expression surrogate(SearchStats searchStats);
+}

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

@@ -9,7 +9,6 @@
 
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
-import org.elasticsearch.xpack.esql.stats.SearchStats;
 
 /**
  * Interface signaling to the planner that the declaring expression
@@ -29,8 +28,4 @@ public interface SurrogateExpression {
      * be replaced.
      */
     Expression surrogate();
-
-    default Expression surrogate(SearchStats searchStats) {
-        return null;
-    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/grouping/Bucket.java

@@ -14,42 +14,33 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.core.TimeValue;
-import org.elasticsearch.logging.LogManager;
-import org.elasticsearch.logging.Logger;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
-import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.expression.Foldables;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
 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.MultiTypeEsField;
-import org.elasticsearch.xpack.esql.expression.SurrogateExpression;
+import org.elasticsearch.xpack.esql.expression.LocalSurrogateExpression;
 import org.elasticsearch.xpack.esql.expression.function.Example;
 import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
 import org.elasticsearch.xpack.esql.expression.function.FunctionType;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.expression.function.TwoOptionalArguments;
 import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Floor;
-import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Div;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Mul;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.stats.SearchStats;
 
 import java.io.IOException;
-import java.time.ZoneId;
-import java.time.ZoneOffset;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
-import java.util.stream.Collectors;
 
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
@@ -58,10 +49,10 @@
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.THIRD;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isNumeric;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
-import static org.elasticsearch.xpack.esql.core.type.DataType.isDateTime;
 import static org.elasticsearch.xpack.esql.expression.Validations.isFoldable;
+import static org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc.maybeSubstituteWithRoundTo;
+import static org.elasticsearch.xpack.esql.session.Configuration.DEFAULT_TZ;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateTimeToLong;
-import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateWithTypeToString;
 
 /**
  * Splits dates and numbers into a given number of buckets. There are two ways to invoke
@@ -73,11 +64,9 @@ public class Bucket extends GroupingFunction.EvaluatableGroupingFunction
     implements
         PostOptimizationVerificationAware,
         TwoOptionalArguments,
-        SurrogateExpression {
+        LocalSurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Bucket", Bucket::new);
 
-    private static final Logger logger = LogManager.getLogger(Bucket.class);
-
     // TODO maybe we should just cover the whole of representable dates here - like ten years, 100 years, 1000 years, all the way up.
     // That way you never end up with more than the target number of buckets.
     private static final Rounding LARGEST_HUMAN_DATE_ROUNDING = Rounding.builder(Rounding.DateTimeUnit.YEAR_OF_CENTURY).build();
@@ -101,8 +90,6 @@ public class Bucket extends GroupingFunction.EvaluatableGroupingFunction
         Rounding.builder(TimeValue.timeValueMillis(10)).build(),
         Rounding.builder(TimeValue.timeValueMillis(1)).build(), };
 
-    private static final ZoneId DEFAULT_TZ = ZoneOffset.UTC; // TODO: plug in the config
-
     private final Expression field;
     private final Expression buckets;
     private final Expression from;
@@ -510,41 +497,16 @@ public String toString() {
         return "Bucket{" + "field=" + field + ", buckets=" + buckets + ", from=" + from + ", to=" + to + '}';
     }
 
-    @Override
-    public Expression surrogate() {
-        return null;
-    }
-
     @Override
     public Expression surrogate(SearchStats searchStats) {
-        if (field() instanceof FieldAttribute fa && fa.field() instanceof MultiTypeEsField == false && isDateTime(fa.dataType())) {
-            // Extract min/max from SearchStats
-            DataType fieldType = fa.dataType();
-            FieldAttribute.FieldName fieldName = fa.fieldName();
-            var min = searchStats.min(fieldName);
-            var max = searchStats.max(fieldName);
-            // If min/max is available create rounding with them
-            if (min instanceof Long minValue && max instanceof Long maxValue && buckets().foldable()) {
-                Rounding.Prepared rounding = getDateRounding(FoldContext.small(), minValue, maxValue);
-                long[] roundingPoints = rounding.fixedRoundingPoints();
-                if (roundingPoints == null) {
-                    logger.trace(
-                        "Fixed rounding point is null for field {}, minValue {} in string format {} and maxValue {} in string format {}",
-                        fieldName,
-                        minValue,
-                        dateWithTypeToString(minValue, fieldType),
-                        maxValue,
-                        dateWithTypeToString(maxValue, fieldType)
-                    );
-                    return null;
-                }
-                // Convert to round_to function with the roundings
-                List<Expression> points = Arrays.stream(roundingPoints)
-                    .mapToObj(l -> new Literal(Source.EMPTY, l, fieldType))
-                    .collect(Collectors.toList());
-                return new RoundTo(source(), field(), points);
-            }
-        }
-        return null;
+        // LocalSubstituteSurrogateExpressions should make sure this doesn't happen
+        assert searchStats != null : "SearchStats cannot be null";
+        return maybeSubstituteWithRoundTo(
+            source(),
+            field(),
+            buckets(),
+            searchStats,
+            (interval, minValue, maxValue) -> getDateRounding(FoldContext.small(), minValue, maxValue)
+        );
     }
 }

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

@@ -57,7 +57,7 @@
 public class DateDiff extends EsqlScalarFunction {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "DateDiff", DateDiff::new);
 
-    public static final ZoneId UTC = ZoneId.of("Z");
+    public static final ZoneId UTC = org.elasticsearch.xpack.esql.core.util.DateUtils.UTC;
 
     private final Expression unit;
     private final Expression startTimestamp;

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

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.expression.function.scalar.date;
 
 import org.elasticsearch.common.Rounding;
+import org.elasticsearch.common.TriFunction;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
@@ -26,7 +27,7 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.MultiTypeEsField;
-import org.elasticsearch.xpack.esql.expression.SurrogateExpression;
+import org.elasticsearch.xpack.esql.expression.LocalSurrogateExpression;
 import org.elasticsearch.xpack.esql.expression.function.Example;
 import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
 import org.elasticsearch.xpack.esql.expression.function.Param;
@@ -39,7 +40,6 @@
 import java.time.Duration;
 import java.time.Period;
 import java.time.ZoneId;
-import java.time.ZoneOffset;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -52,9 +52,10 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.DATETIME;
 import static org.elasticsearch.xpack.esql.core.type.DataType.DATE_NANOS;
 import static org.elasticsearch.xpack.esql.core.type.DataType.isDateTime;
+import static org.elasticsearch.xpack.esql.session.Configuration.DEFAULT_TZ;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateWithTypeToString;
 
-public class DateTrunc extends EsqlScalarFunction implements SurrogateExpression {
+public class DateTrunc extends EsqlScalarFunction implements LocalSurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         Expression.class,
         "DateTrunc",
@@ -74,7 +75,6 @@ public interface DateTruncFactoryProvider {
     );
     private final Expression interval;
     private final Expression timestampField;
-    protected static final ZoneId DEFAULT_TZ = ZoneOffset.UTC;
 
     @FunctionInfo(
         returnType = { "date", "date_nanos" },
@@ -287,22 +287,35 @@ public static ExpressionEvaluator.Factory evaluator(
     }
 
     @Override
-    public Expression surrogate() { // there is no substitute without SearchStats
-        return null;
+    public Expression surrogate(SearchStats searchStats) {
+        // LocalSubstituteSurrogateExpressions should make sure this doesn't happen
+        assert searchStats != null : "SearchStats cannot be null";
+        return maybeSubstituteWithRoundTo(
+            source(),
+            field(),
+            interval(),
+            searchStats,
+            (interval, minValue, maxValue) -> createRounding(interval, DEFAULT_TZ, minValue, maxValue)
+        );
     }
 
-    @Override
-    public Expression surrogate(SearchStats searchStats) {
-        if (field() instanceof FieldAttribute fa && fa.field() instanceof MultiTypeEsField == false && isDateTime(fa.dataType())) {
+    public static RoundTo maybeSubstituteWithRoundTo(
+        Source source,
+        Expression field,
+        Expression foldableTimeExpression,
+        SearchStats searchStats,
+        TriFunction<Object, Long, Long, Rounding.Prepared> roundingFunction
+    ) {
+        if (field instanceof FieldAttribute fa && fa.field() instanceof MultiTypeEsField == false && isDateTime(fa.dataType())) {
             // Extract min/max from SearchStats
             DataType fieldType = fa.dataType();
             FieldAttribute.FieldName fieldName = fa.fieldName();
             var min = searchStats.min(fieldName);
             var max = searchStats.max(fieldName);
             // If min/max is available create rounding with them
-            if (min instanceof Long minValue && max instanceof Long maxValue && interval().foldable()) {
-                Object foldedInterval = interval().fold(FoldContext.small() /* TODO remove me */);
-                Rounding.Prepared rounding = createRounding(foldedInterval, DEFAULT_TZ, minValue, maxValue);
+            if (min instanceof Long minValue && max instanceof Long maxValue && foldableTimeExpression.foldable()) {
+                Object foldedInterval = foldableTimeExpression.fold(FoldContext.small() /* TODO remove me */);
+                Rounding.Prepared rounding = roundingFunction.apply(foldedInterval, minValue, maxValue);
                 long[] roundingPoints = rounding.fixedRoundingPoints();
                 if (roundingPoints == null) {
                     logger.trace(
@@ -319,7 +332,7 @@ public Expression surrogate(SearchStats searchStats) {
                 List<Expression> points = Arrays.stream(roundingPoints)
                     .mapToObj(l -> new Literal(Source.EMPTY, l, fieldType))
                     .collect(Collectors.toList());
-                return new RoundTo(source(), field(), points);
+                return new RoundTo(source, field, points);
             }
         }
         return null;

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

@@ -15,9 +15,9 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.InferIsNotNull;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.InferNonNullAggConstraint;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.LocalPropagateEmptyRelation;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.LocalSubstituteSurrogateExpressions;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.ReplaceFieldWithConstantOrNull;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.ReplaceTopNWithLimitAndSort;
-import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.SubstituteSurrogateExpressionsWithSearchStats;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.rule.ParameterizedRuleExecutor;
 import org.elasticsearch.xpack.esql.rule.Rule;
@@ -48,7 +48,7 @@ public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<Logical
             new ReplaceFieldWithConstantOrNull(),
             new InferIsNotNull(),
             new InferNonNullAggConstraint(),
-            new SubstituteSurrogateExpressionsWithSearchStats()
+            new LocalSubstituteSurrogateExpressions()
         ),
         localOperators(),
         cleanup()