ReplaceDateTruncBucketWithRoundTo

Author: fang-xing-esql

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

@@ -25,7 +25,6 @@
 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.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;
@@ -35,9 +34,7 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Floor;
 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.expression.predicate.operator.comparison.EsqlBinaryComparison;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
-import org.elasticsearch.xpack.esql.stats.SearchStats;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -51,7 +48,6 @@
 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.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;
 
@@ -64,8 +60,7 @@
 public class Bucket extends GroupingFunction.EvaluatableGroupingFunction
     implements
         PostOptimizationVerificationAware,
-        TwoOptionalArguments,
-        LocalSurrogateExpression {
+        TwoOptionalArguments {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Bucket", Bucket::new);
 
     // TODO maybe we should just cover the whole of representable dates here - like ten years, 100 years, 1000 years, all the way up.
@@ -306,7 +301,7 @@ private Rounding.Prepared getDateRounding(FoldContext foldContext) {
         return getDateRounding(foldContext, null, null);
     }
 
-    private Rounding.Prepared getDateRounding(FoldContext foldContext, Long min, Long max) {
+    public Rounding.Prepared getDateRounding(FoldContext foldContext, Long min, Long max) {
         assert field.dataType() == DataType.DATETIME || field.dataType() == DataType.DATE_NANOS : "expected date type; got " + field;
         if (buckets.dataType().isWholeNumber()) {
             int b = ((Number) buckets.fold(foldContext)).intValue();
@@ -497,18 +492,4 @@ public Expression to() {
     public String toString() {
         return "Bucket{" + "field=" + field + ", buckets=" + buckets + ", from=" + from + ", to=" + to + '}';
     }
-
-    @Override
-    public Expression surrogate(SearchStats searchStats, List<EsqlBinaryComparison> binaryComparisons) {
-        // LocalSubstituteSurrogateExpressions should make sure this doesn't happen
-        assert searchStats != null : "SearchStats cannot be null";
-        return maybeSubstituteWithRoundTo(
-            source(),
-            field(),
-            buckets(),
-            searchStats,
-            binaryComparisons,
-            (interval, minValue, maxValue) -> getDateRounding(FoldContext.small(), minValue, maxValue)
-        );
-    }
 }

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

@@ -8,7 +8,6 @@
 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;
@@ -17,61 +16,38 @@
 import org.elasticsearch.compute.ann.Fixed;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.core.TimeValue;
-import org.elasticsearch.core.Tuple;
-import org.elasticsearch.logging.LogManager;
-import org.elasticsearch.logging.Logger;
 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.Literal;
 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.core.util.Holder;
-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;
 import org.elasticsearch.xpack.esql.expression.function.scalar.EsqlScalarFunction;
-import org.elasticsearch.xpack.esql.expression.function.scalar.math.RoundTo;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.GreaterThan;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.GreaterThanOrEqual;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThan;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
-import org.elasticsearch.xpack.esql.stats.SearchStats;
 
 import java.io.IOException;
 import java.time.Duration;
 import java.time.Period;
 import java.time.ZoneId;
-import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
 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 LocalSurrogateExpression {
+public class DateTrunc extends EsqlScalarFunction {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         Expression.class,
         "DateTrunc",
         DateTrunc::new
     );
 
-    private static final Logger logger = LogManager.getLogger(DateTrunc.class);
-
     @FunctionalInterface
     public interface DateTruncFactoryProvider {
         ExpressionEvaluator.Factory apply(Source source, ExpressionEvaluator.Factory lhs, Rounding.Prepared rounding);
@@ -129,7 +105,7 @@ public String getWriteableName() {
         return ENTRY.name;
     }
 
-    Expression interval() {
+    public Expression interval() {
         return interval;
     }
 
@@ -293,96 +269,4 @@ public static ExpressionEvaluator.Factory evaluator(
     ) {
         return evaluatorMap.get(forType).apply(source, fieldEvaluator, rounding);
     }
-
-    @Override
-    public Expression surrogate(SearchStats searchStats, List<EsqlBinaryComparison> binaryComparisons) {
-        // LocalSubstituteSurrogateExpressions should make sure this doesn't happen
-        assert searchStats != null : "SearchStats cannot be null";
-        return maybeSubstituteWithRoundTo(
-            source(),
-            field(),
-            interval(),
-            searchStats,
-            binaryComparisons,
-            (interval, minValue, maxValue) -> createRounding(interval, DEFAULT_TZ, minValue, maxValue)
-        );
-    }
-
-    public static RoundTo maybeSubstituteWithRoundTo(
-        Source source,
-        Expression field,
-        Expression foldableTimeExpression,
-        SearchStats searchStats,
-        List<EsqlBinaryComparison> binaryComparisons,
-        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);
-            // Extract min/max from query
-            Tuple<Long, Long> minMaxFromPredicates = minMaxFromPredicates(binaryComparisons);
-            Long minFromPredicates = minMaxFromPredicates.v1();
-            Long maxFromPredicates = minMaxFromPredicates.v2();
-            // Consolidate min/max from SearchStats and query
-            if (minFromPredicates instanceof Long minValue) {
-                min = min instanceof Long m ? Math.max(m, minValue) : minValue;
-            }
-            if (maxFromPredicates instanceof Long maxValue) {
-                max = max instanceof Long m ? Math.min(m, maxValue) : maxValue;
-            }
-            // If min/max is available create rounding with them
-            if (min instanceof Long minValue && max instanceof Long maxValue && foldableTimeExpression.foldable() && minValue <= maxValue) {
-                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(
-                        "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;
-    }
-
-    private static Tuple<Long, Long> minMaxFromPredicates(List<EsqlBinaryComparison> binaryComparisons) {
-        long[] min = new long[] { Long.MIN_VALUE };
-        long[] max = new long[] { Long.MAX_VALUE };
-        Holder<Boolean> foundMinValue = new Holder<>(false);
-        Holder<Boolean> foundMaxValue = new Holder<>(false);
-        for (EsqlBinaryComparison binaryComparison : binaryComparisons) {
-            if (binaryComparison.right() instanceof Literal l) {
-                long value = Long.parseLong(l.value().toString());
-                if (binaryComparison instanceof Equals) {
-                    return new Tuple<>(value, value);
-                }
-                if (binaryComparison instanceof GreaterThan || binaryComparison instanceof GreaterThanOrEqual) {
-                    if (value >= min[0]) {
-                        min[0] = value;
-                        foundMinValue.set(true);
-                    }
-                } else if (binaryComparison instanceof LessThan || binaryComparison instanceof LessThanOrEqual) {
-                    if (value <= max[0]) {
-                        max[0] = value;
-                        foundMaxValue.set(true);
-                    }
-                }
-            }
-        }
-        return new Tuple<>(foundMinValue.get() ? min[0] : null, foundMaxValue.get() ? max[0] : null);
-    }
 }

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

@@ -15,7 +15,7 @@
 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.ReplaceDateTruncBucketWithRoundTo;
 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.plan.logical.LogicalPlan;
@@ -48,7 +48,7 @@ public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<Logical
             new ReplaceFieldWithConstantOrNull(),
             new InferIsNotNull(),
             new InferNonNullAggConstraint(),
-            new LocalSubstituteSurrogateExpressions()
+            new ReplaceDateTruncBucketWithRoundTo()
         ),
         localOperators(),
         cleanup()