Add "emitEmptyBuckets" parameter to the "Bucket" function.

Author: przemekwitek

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

@@ -301,7 +301,7 @@ private static FunctionDefinition[][] functions() {
         return new FunctionDefinition[][] {
             // grouping functions
             new FunctionDefinition[] {
-                def(Bucket.class, Bucket::new, "bucket", "bin"),
+                def(Bucket.class, quin(Bucket::new), "bucket", "bin"),
                 def(Categorize.class, Categorize::new, "categorize") },
             // aggregate functions
             // since they declare two public constructors - one with filter (for nested where) and one without
@@ -1010,6 +1010,39 @@ protected static <T extends Function> FunctionDefinition def(Class<T> function,
         return def(function, builder, names);
     }
 
+    /**
+     * Build a {@linkplain FunctionDefinition} for a quinary function.
+     */
+    @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, QuinaryBuilder<T> ctorRef, String... names) {
+        FunctionBuilder builder = (source, children, cfg) -> {
+            if (OptionalArgument.class.isAssignableFrom(function)) {
+                if (children.size() > 5 || children.size() < 4) {
+                    throw new QlIllegalArgumentException("expects four or five arguments");
+                }
+            } else if (TwoOptionalArguments.class.isAssignableFrom(function)) {
+                if (children.size() > 5 || children.size() < 3) {
+                    throw new QlIllegalArgumentException("expects minimum three, maximum five arguments");
+                }
+            } else if (ThreeOptionalArguments.class.isAssignableFrom(function)) {
+                if (children.size() > 5 || children.size() < 2) {
+                    throw new QlIllegalArgumentException("expects minimum two, maximum five arguments");
+                }
+            } else if (children.size() != 5) {
+                throw new QlIllegalArgumentException("expects exactly five 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,
+                children.size() > 4 ? children.get(4) : null
+            );
+        };
+        return def(function, builder, names);
+    }
+
     protected interface QuaternaryBuilder<T> {
         T build(Source source, Expression one, Expression two, Expression three, Expression four);
     }
@@ -1204,4 +1237,7 @@ private static <T extends Function> TernaryBuilder<T> tri(TernaryBuilder<T> func
         return function;
     }
 
+    private static <T extends Function> QuinaryBuilder<T> quin(QuinaryBuilder<T> function) {
+        return function;
+    }
 }

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

@@ -0,0 +1,16 @@
+/*
+ * 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.function;
+
+/**
+ * Marker interface indicating that a function accepts three optional arguments (the last three).
+ * This is used by the {@link EsqlFunctionRegistry} to perform validation of function declaration.
+ */
+public interface ThreeOptionalArguments {
+
+}

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

@@ -29,7 +29,7 @@
 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.ThreeOptionalArguments;
 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.predicate.operator.arithmetic.Div;
@@ -61,7 +61,7 @@
 public class Bucket extends GroupingFunction.EvaluatableGroupingFunction
     implements
         PostOptimizationVerificationAware,
-        TwoOptionalArguments {
+        ThreeOptionalArguments {
     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.
@@ -93,6 +93,7 @@ public class Bucket extends GroupingFunction.EvaluatableGroupingFunction
     private final Expression buckets;
     private final Expression from;
     private final Expression to;
+    private final Expression emitEmptyBuckets;
 
     @FunctionInfo(
         returnType = { "double", "date", "date_nanos" },
@@ -211,13 +212,20 @@ 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,
+        @Param(
+            name = "emitEmptyBuckets",
+            type = { "boolean" },
+            optional = true,
+            description = "Whether or not empty buckets should be emitted."
+        ) Expression emitEmptyBuckets
     ) {
-        super(source, fields(field, buckets, from, to));
+        super(source, fields(field, buckets, from, to, emitEmptyBuckets));
         this.field = field;
         this.buckets = buckets;
         this.from = from;
         this.to = to;
+        this.emitEmptyBuckets = emitEmptyBuckets;
     }
 
     private Bucket(StreamInput in) throws IOException {
@@ -226,11 +234,12 @@ 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)
         );
     }
 
-    private static List<Expression> fields(Expression field, Expression buckets, Expression from, Expression to) {
+    private static List<Expression> fields(Expression field, Expression buckets, Expression from, Expression to, Expression emitEmptyBuckets) {
         List<Expression> list = new ArrayList<>(4);
         list.add(field);
         list.add(buckets);
@@ -240,6 +249,9 @@ private static List<Expression> fields(Expression field, Expression buckets, Exp
                 list.add(to);
             }
         }
+        if (emitEmptyBuckets != null) {
+            list.add(emitEmptyBuckets);
+        }
         return list;
     }
 
@@ -250,6 +262,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeNamedWriteable(buckets);
         out.writeOptionalNamedWriteable(from);
         out.writeOptionalNamedWriteable(to);
+        out.writeOptionalNamedWriteable(emitEmptyBuckets);
     }
 
     @Override
@@ -259,7 +272,7 @@ public String getWriteableName() {
 
     @Override
     public boolean foldable() {
-        return field.foldable() && buckets.foldable() && (from == null || from.foldable()) && (to == null || to.foldable());
+        return field.foldable() && buckets.foldable() && (from == null || from.foldable()) && (to == null || to.foldable()) && (emitEmptyBuckets == null || emitEmptyBuckets.foldable());
     }
 
     @Override
@@ -460,12 +473,13 @@ 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);
+        Expression emitEmptyBuckets = newChildren.size() > 4 ? newChildren.get(4) : null;
+        return new Bucket(source(), newChildren.get(0), newChildren.get(1), from, to, emitEmptyBuckets);
     }
 
     @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, emitEmptyBuckets);
     }
 
     public Expression field() {
@@ -484,8 +498,12 @@ public Expression to() {
         return to;
     }
 
+    public Expression emitEmptyBuckets() {
+        return emitEmptyBuckets;
+    }
+
     @Override
     public String toString() {
-        return "Bucket{" + "field=" + field + ", buckets=" + buckets + ", from=" + from + ", to=" + to + '}';
+        return "Bucket{" + "field=" + field + ", buckets=" + buckets + ", from=" + from + ", to=" + to + ", emitEmptyBuckets=" + emitEmptyBuckets + '}';
     }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/RailRoadDiagram.java

@@ -72,7 +72,10 @@ static String functionSignature(FunctionDefinition definition) throws IOExceptio
                             // BUCKET requires optional args to be optional together, so we need custom code to do that
                             var nextArg = args.get(++i);
                             assert nextArg.optional();
-                            Sequence seq = new Sequence(new Literal(argName), new Syntax(","), new Literal(nextArg.name));
+                            var nexterArg = args.get(++i);
+                            assert nexterArg.optional();
+                            // TODO: Should it be possible to be able to specify "emitEmptyBuckets" but not "from" and "to"?
+                            Sequence seq = new Sequence(new Literal(argName), new Syntax(","), new Literal(nextArg.name), new Syntax(","), new Literal(nexterArg.name));
                             argExpressions.add(new Repetition(seq, 0, 1));
                         } else if (i < args.size() - 1 && args.get(i + 1).optional() == false) {
                             // Special case with leading optional args

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/grouping/BucketSerializationTests.java

@@ -21,7 +21,8 @@ protected Bucket createTestInstance() {
         Expression buckets = randomChild();
         Expression from = randomChild();
         Expression to = randomChild();
-        return new Bucket(source, field, buckets, from, to);
+        Expression emitEmptyBuckets = randomChild();
+        return new Bucket(source, field, buckets, from, to, emitEmptyBuckets);
     }
 
     @Override
@@ -31,12 +32,14 @@ protected Bucket mutateInstance(Bucket instance) throws IOException {
         Expression buckets = instance.buckets();
         Expression from = instance.from();
         Expression to = instance.to();
-        switch (between(0, 3)) {
+        Expression emitEmptyBuckets = instance.emitEmptyBuckets();
+        switch (between(0, 4)) {
             case 0 -> field = randomValueOtherThan(field, AbstractExpressionSerializationTests::randomChild);
             case 1 -> buckets = randomValueOtherThan(buckets, AbstractExpressionSerializationTests::randomChild);
             case 2 -> from = randomValueOtherThan(from, AbstractExpressionSerializationTests::randomChild);
             case 3 -> to = randomValueOtherThan(to, AbstractExpressionSerializationTests::randomChild);
+            case 4 -> emitEmptyBuckets = randomValueOtherThan(emitEmptyBuckets, AbstractExpressionSerializationTests::randomChild);
         }
-        return new Bucket(source, field, buckets, from, to);
+        return new Bucket(source, field, buckets, from, to, emitEmptyBuckets);
     }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/grouping/BucketTests.java

@@ -319,10 +319,14 @@ private static Matcher<Object> resultsMatcher(List<TestCaseSupplier.TypedData> t
     protected Expression build(Source source, List<Expression> args) {
         Expression from = null;
         Expression to = null;
+        Expression emitEmptyBuckets = null;
         if (args.size() > 2) {
             from = args.get(2);
             to = args.get(3);
         }
-        return new Bucket(source, args.get(0), args.get(1), from, to);
+        if (args.size() > 4) {
+            emitEmptyBuckets = args.get(4);
+        }
+        return new Bucket(source, args.get(0), args.get(1), from, to, emitEmptyBuckets);
     }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/FoldNullTests.java

@@ -265,7 +265,7 @@ public void testNullFoldableDoesNotApplyToIsNullAndNotNull() {
     }
 
     public void testNullBucketGetsFolded() {
-        assertEquals(NULL, foldNull(new Bucket(EMPTY, NULL, NULL, NULL, NULL)));
+        assertEquals(NULL, foldNull(new Bucket(EMPTY, NULL, NULL, NULL, NULL, NULL)));
     }
 
     public void testNullCategorizeGroupingNotFolded() {