Add configuration to DATE_TRUNC and related functions

Author: ivancea

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

@@ -344,7 +344,7 @@ private static FunctionDefinition[][] functions() {
             new FunctionDefinition[] {
                 def(Bucket.class, Bucket::new, "bucket", "bin"),
                 def(Categorize.class, Categorize::new, "categorize"),
-                def(TBucket.class, uni(TBucket::new), "tbucket") },
+                def(TBucket.class, (UnaryConfigurationAwareBuilder<TBucket>) TBucket::new, "tbucket") },
             // aggregate functions
             // since they declare two public constructors - one with filter (for nested where) and one without
             // use casting to disambiguate between the two
@@ -1276,6 +1276,43 @@ protected interface TernaryConfigurationAwareBuilder<T> {
         T build(Source source, Expression one, Expression two, Expression three, Configuration configuration);
     }
 
+    /**
+     * Build a {@linkplain FunctionDefinition} for a quaternary function that is configuration aware.
+     */
+    @SuppressWarnings("overloads")  // These are ambiguous if you aren't using ctor references but we always do
+    protected static <T extends Function> FunctionDefinition def(
+        Class<T> function,
+        QuaternaryConfigurationAwareBuilder<T> ctorRef,
+        String... names
+    ) {
+        FunctionBuilder builder = (source, children, cfg) -> {
+            if (OptionalArgument.class.isAssignableFrom(function)) {
+                if (children.size() > 4 || children.size() < 3) {
+                    throw new QlIllegalArgumentException("expects three or four arguments");
+                }
+            } else if (TwoOptionalArguments.class.isAssignableFrom(function)) {
+                if (children.size() > 4 || children.size() < 2) {
+                    throw new QlIllegalArgumentException("expects minimum two, maximum four arguments");
+                }
+            } else if (children.size() != 4) {
+                throw new QlIllegalArgumentException("expects exactly four arguments");
+            }
+            return ctorRef.build(
+                source,
+                children.get(0),
+                children.get(1),
+                children.size() > 2 ? children.get(2) : null,
+                children.size() > 3 ? children.get(3) : null,
+                cfg
+            );
+        };
+        return def(function, builder, names);
+    }
+
+    protected interface QuaternaryConfigurationAwareBuilder<T> {
+        T build(Source source, Expression one, Expression two, Expression three, Expression four, Configuration configuration);
+    }
+
     //
     // Utility functions to help disambiguate the method handle passed in.
     // They work by providing additional method information to help the compiler know which method to pick.

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

@@ -35,10 +35,13 @@
 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.session.Configuration;
 
 import java.io.IOException;
+import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
@@ -48,7 +51,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.session.Configuration.DEFAULT_TZ;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateTimeToLong;
 
 /**
@@ -65,27 +67,34 @@ public class Bucket extends GroupingFunction.EvaluatableGroupingFunction
 
     // 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();
-    private static final Rounding[] HUMAN_DATE_ROUNDINGS = new Rounding[] {
-        Rounding.builder(Rounding.DateTimeUnit.MONTH_OF_YEAR).build(),
-        Rounding.builder(Rounding.DateTimeUnit.WEEK_OF_WEEKYEAR).build(),
-        Rounding.builder(Rounding.DateTimeUnit.DAY_OF_MONTH).build(),
-        Rounding.builder(TimeValue.timeValueHours(12)).build(),
-        Rounding.builder(TimeValue.timeValueHours(3)).build(),
-        Rounding.builder(TimeValue.timeValueHours(1)).build(),
-        Rounding.builder(TimeValue.timeValueMinutes(30)).build(),
-        Rounding.builder(TimeValue.timeValueMinutes(10)).build(),
-        Rounding.builder(TimeValue.timeValueMinutes(5)).build(),
-        Rounding.builder(TimeValue.timeValueMinutes(1)).build(),
-        Rounding.builder(TimeValue.timeValueSeconds(30)).build(),
-        Rounding.builder(TimeValue.timeValueSeconds(10)).build(),
-        Rounding.builder(TimeValue.timeValueSeconds(5)).build(),
-        Rounding.builder(TimeValue.timeValueSeconds(1)).build(),
-        Rounding.builder(TimeValue.timeValueMillis(100)).build(),
-        Rounding.builder(TimeValue.timeValueMillis(50)).build(),
-        Rounding.builder(TimeValue.timeValueMillis(10)).build(),
-        Rounding.builder(TimeValue.timeValueMillis(1)).build(), };
-
+    private static final Rounding.Builder LARGEST_HUMAN_DATE_ROUNDING = Rounding.builder(Rounding.DateTimeUnit.YEAR_OF_CENTURY);
+    private static final Rounding.Builder[] HUMAN_DATE_ROUNDINGS = new Rounding.Builder[] {
+        Rounding.builder(Rounding.DateTimeUnit.MONTH_OF_YEAR),
+        Rounding.builder(Rounding.DateTimeUnit.WEEK_OF_WEEKYEAR),
+        Rounding.builder(Rounding.DateTimeUnit.DAY_OF_MONTH),
+        Rounding.builder(TimeValue.timeValueHours(12)),
+        Rounding.builder(TimeValue.timeValueHours(3)),
+        Rounding.builder(TimeValue.timeValueHours(1)),
+        Rounding.builder(TimeValue.timeValueMinutes(30)),
+        Rounding.builder(TimeValue.timeValueMinutes(10)),
+        Rounding.builder(TimeValue.timeValueMinutes(5)),
+        Rounding.builder(TimeValue.timeValueMinutes(1)),
+        Rounding.builder(TimeValue.timeValueSeconds(30)),
+        Rounding.builder(TimeValue.timeValueSeconds(10)),
+        Rounding.builder(TimeValue.timeValueSeconds(5)),
+        Rounding.builder(TimeValue.timeValueSeconds(1)),
+        Rounding.builder(TimeValue.timeValueMillis(100)),
+        Rounding.builder(TimeValue.timeValueMillis(50)),
+        Rounding.builder(TimeValue.timeValueMillis(10)),
+        Rounding.builder(TimeValue.timeValueMillis(1)), };
+
+    /*
+     * As Bucket already extends GroupingFunction, it can't extend EsqlConfigurationFunction, so we had to replicate here:
+     * - The Configuration field
+     * - HashCode
+     * - Equals
+     */
+    private final Configuration configuration;
     private final Expression field;
     private final Expression buckets;
     private final Expression from;
@@ -208,13 +217,15 @@ public Bucket(
             type = { "integer", "long", "double", "date", "keyword", "text" },
             optional = true,
             description = "End of the range. Can be a number, a date or a date expressed as a string."
-        ) Expression to
+        ) Expression to,
+        Configuration configuration
     ) {
         super(source, fields(field, buckets, from, to));
         this.field = field;
         this.buckets = buckets;
         this.from = from;
         this.to = to;
+        this.configuration = configuration;
     }
 
     private Bucket(StreamInput in) throws IOException {
@@ -223,10 +234,15 @@ private Bucket(StreamInput in) throws IOException {
             in.readNamedWriteable(Expression.class),
             in.readNamedWriteable(Expression.class),
             in.readOptionalNamedWriteable(Expression.class),
-            in.readOptionalNamedWriteable(Expression.class)
+            in.readOptionalNamedWriteable(Expression.class),
+            ((PlanStreamInput) in).configuration()
         );
     }
 
+    public Configuration configuration() {
+        return configuration;
+    }
+
     private static List<Expression> fields(Expression field, Expression buckets, Expression from, Expression to) {
         List<Expression> list = new ArrayList<>(4);
         list.add(field);
@@ -308,19 +324,20 @@ public Rounding.Prepared getDateRounding(FoldContext foldContext, Long min, Long
             long f = foldToLong(foldContext, from);
             long t = foldToLong(foldContext, to);
             if (min != null && max != null) {
-                return new DateRoundingPicker(b, f, t).pickRounding().prepare(min, max);
+                return new DateRoundingPicker(b, f, t, configuration.zoneId()).pickRounding().prepare(min, max);
             }
-            return new DateRoundingPicker(b, f, t).pickRounding().prepareForUnknown();
+            return new DateRoundingPicker(b, f, t, configuration.zoneId()).pickRounding().prepareForUnknown();
         } else {
             assert DataType.isTemporalAmount(buckets.dataType()) : "Unexpected span data type [" + buckets.dataType() + "]";
-            return DateTrunc.createRounding(buckets.fold(foldContext), DEFAULT_TZ, min, max);
+            return DateTrunc.createRounding(buckets.fold(foldContext), configuration.zoneId(), min, max);
         }
     }
 
-    private record DateRoundingPicker(int buckets, long from, long to) {
+    private record DateRoundingPicker(int buckets, long from, long to, ZoneId zoneId) {
         Rounding pickRounding() {
-            Rounding prev = LARGEST_HUMAN_DATE_ROUNDING;
-            for (Rounding r : HUMAN_DATE_ROUNDINGS) {
+            Rounding prev = LARGEST_HUMAN_DATE_ROUNDING.timeZone(zoneId).build();
+            for (Rounding.Builder builder : HUMAN_DATE_ROUNDINGS) {
+                Rounding r = builder.timeZone(zoneId).build();
                 if (roundingIsOk(r)) {
                     prev = r;
                 } else {
@@ -464,12 +481,12 @@ public DataType dataType() {
     public Expression replaceChildren(List<Expression> newChildren) {
         Expression from = newChildren.size() > 2 ? newChildren.get(2) : null;
         Expression to = newChildren.size() > 3 ? newChildren.get(3) : null;
-        return new Bucket(source(), newChildren.get(0), newChildren.get(1), from, to);
+        return new Bucket(source(), newChildren.get(0), newChildren.get(1), from, to, configuration);
     }
 
     @Override
     protected NodeInfo<? extends Expression> info() {
-        return NodeInfo.create(this, Bucket::new, field, buckets, from, to);
+        return NodeInfo.create(this, Bucket::new, field, buckets, from, to, configuration);
     }
 
     public Expression field() {
@@ -492,4 +509,19 @@ public Expression to() {
     public String toString() {
         return "Bucket{" + "field=" + field + ", buckets=" + buckets + ", from=" + from + ", to=" + to + '}';
     }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(getClass(), children(), configuration);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (super.equals(obj) == false) {
+            return false;
+        }
+        Bucket other = (Bucket) obj;
+
+        return configuration.equals(other.configuration);
+    }
 }

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

@@ -19,9 +19,11 @@
 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.session.Configuration;
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Objects;
 
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
@@ -33,6 +35,13 @@
 public class TBucket extends GroupingFunction.EvaluatableGroupingFunction implements SurrogateExpression {
     public static final String NAME = "TBucket";
 
+    /*
+     * As Bucket already extends GroupingFunction, it can't extend EsqlConfigurationFunction, so we had to replicate here:
+     * - The Configuration field
+     * - HashCode
+     * - Equals
+     */
+    private final Configuration configuration;
     private final Expression buckets;
     private final Expression timestamp;
 
@@ -62,19 +71,22 @@ public class TBucket extends GroupingFunction.EvaluatableGroupingFunction implem
     )
     public TBucket(
         Source source,
-        @Param(name = "buckets", type = { "date_period", "time_duration" }, description = "Desired bucket size.") Expression buckets
+        @Param(name = "buckets", type = { "date_period", "time_duration" }, description = "Desired bucket size.") Expression buckets,
+        Configuration configuration
     ) {
         this(
             source,
             buckets,
-            new UnresolvedTimestamp(source, "TBucket function requires @timestamp field, but @timestamp was renamed or dropped")
+            new UnresolvedTimestamp(source, "TBucket function requires @timestamp field, but @timestamp was renamed or dropped"),
+            configuration
         );
     }
 
-    public TBucket(Source source, Expression buckets, Expression timestamp) {
+    public TBucket(Source source, Expression buckets, Expression timestamp, Configuration configuration) {
         super(source, List.of(buckets, timestamp));
         this.buckets = buckets;
         this.timestamp = timestamp;
+        this.configuration = configuration;
     }
 
     @Override
@@ -94,7 +106,7 @@ public EvalOperator.ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvalua
 
     @Override
     public Expression surrogate() {
-        return new Bucket(source(), timestamp, buckets, null, null);
+        return new Bucket(source(), timestamp, buckets, null, null, configuration);
     }
 
     @Override
@@ -114,12 +126,12 @@ public DataType dataType() {
 
     @Override
     public Expression replaceChildren(List<Expression> newChildren) {
-        return new TBucket(source(), newChildren.get(0), newChildren.get(1));
+        return new TBucket(source(), newChildren.get(0), newChildren.get(1), configuration);
     }
 
     @Override
     protected NodeInfo<? extends Expression> info() {
-        return NodeInfo.create(this, TBucket::new, buckets, timestamp);
+        return NodeInfo.create(this, TBucket::new, buckets, timestamp, configuration);
     }
 
     public Expression field() {
@@ -134,4 +146,19 @@ public Expression buckets() {
     public String toString() {
         return "TBucket{buckets=" + buckets + "}";
     }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(getClass(), children(), configuration);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (super.equals(obj) == false) {
+            return false;
+        }
+        TBucket other = (TBucket) obj;
+
+        return configuration.equals(other.configuration);
+    }
 }