ESQL: Move more function serialization (#109950)

This moves yet more of the function serialization code from `PlanNamedTypes` into `NamedWriteable` to line up better with the rest of Elasticsearch. This moves 13 more functions including all "configuration" functions and all date functions and all varags functions. There are 38 left.
elastic · Jun 20, 2024 · 7b81e0c · 7b81e0c
1 parent 4a0b026
commit 7b81e0c
Show file tree

Hide file tree

Showing 39 changed files with 1,143 additions and 156 deletions.
diff --git a/...ava/org/elasticsearch/xpack/esql/core/expression/function/scalar/UnaryScalarFunction.java b/...ava/org/elasticsearch/xpack/esql/core/expression/function/scalar/UnaryScalarFunction.java
@@ -6,10 +6,15 @@
  */
 package org.elasticsearch.xpack.esql.core.expression.function.scalar;
 
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.gen.processor.Processor;
 import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamInput;
+import org.elasticsearch.xpack.esql.core.util.PlanStreamOutput;
 
+import java.io.IOException;
 import java.util.List;
 
 import static java.util.Collections.singletonList;
@@ -18,16 +23,21 @@ public abstract class UnaryScalarFunction extends ScalarFunction {
 
     private final Expression field;
 
-    protected UnaryScalarFunction(Source source) {
-        super(source);
-        this.field = null;
-    }
-
     protected UnaryScalarFunction(Source source, Expression field) {
         super(source, singletonList(field));
         this.field = field;
     }
 
+    protected UnaryScalarFunction(StreamInput in) throws IOException {
+        this(Source.readFrom((StreamInput & PlanStreamInput) in), ((PlanStreamInput) in).readExpression());
+    }
+
+    @Override
+    public final void writeTo(StreamOutput out) throws IOException {
+        source().writeTo(out);
+        ((PlanStreamOutput) out).writeExpression(field);
+    }
+
     @Override
     public final UnaryScalarFunction replaceChildren(List<Expression> newChildren) {
         return replaceChild(newChildren.get(0));

diff --git a/...ore/src/main/java/org/elasticsearch/xpack/esql/core/expression/predicate/logical/Not.java b/...ore/src/main/java/org/elasticsearch/xpack/esql/core/expression/predicate/logical/Not.java
@@ -6,6 +6,8 @@
  */
 package org.elasticsearch.xpack.esql.core.expression.predicate.logical;
 
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.function.scalar.UnaryScalarFunction;
 import org.elasticsearch.xpack.esql.core.expression.gen.processor.Processor;
@@ -14,15 +16,27 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 
+import java.io.IOException;
+
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isBoolean;
 
 public class Not extends UnaryScalarFunction implements Negatable<Expression> {
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Not", Not::new);
 
     public Not(Source source, Expression child) {
         super(source, child);
     }
 
+    private Not(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
     @Override
     protected NodeInfo<Not> info() {
         return NodeInfo.create(this, Not::new, field());

diff --git a/...src/main/java/org/elasticsearch/xpack/esql/core/expression/predicate/nulls/IsNotNull.java b/...src/main/java/org/elasticsearch/xpack/esql/core/expression/predicate/nulls/IsNotNull.java
@@ -6,6 +6,8 @@
  */
 package org.elasticsearch.xpack.esql.core.expression.predicate.nulls;
 
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
 import org.elasticsearch.xpack.esql.core.expression.function.scalar.UnaryScalarFunction;
@@ -16,12 +18,28 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 
+import java.io.IOException;
+
 public class IsNotNull extends UnaryScalarFunction implements Negatable<UnaryScalarFunction> {
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
+        Expression.class,
+        "IsNotNull",
+        IsNotNull::new
+    );
 
     public IsNotNull(Source source, Expression field) {
         super(source, field);
     }
 
+    private IsNotNull(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
     @Override
     protected NodeInfo<IsNotNull> info() {
         return NodeInfo.create(this, IsNotNull::new, field());

diff --git a/...re/src/main/java/org/elasticsearch/xpack/esql/core/expression/predicate/nulls/IsNull.java b/...re/src/main/java/org/elasticsearch/xpack/esql/core/expression/predicate/nulls/IsNull.java
@@ -6,6 +6,8 @@
  */
 package org.elasticsearch.xpack.esql.core.expression.predicate.nulls;
 
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
 import org.elasticsearch.xpack.esql.core.expression.function.scalar.UnaryScalarFunction;
@@ -16,12 +18,24 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 
+import java.io.IOException;
+
 public class IsNull extends UnaryScalarFunction implements Negatable<UnaryScalarFunction> {
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "IsNull", IsNull::new);
 
     public IsNull(Source source, Expression field) {
         super(source, field);
     }
 
+    private IsNull(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
     @Override
     protected NodeInfo<IsNull> info() {
         return NodeInfo.create(this, IsNull::new, field());

diff --git a/...main/java/org/elasticsearch/xpack/esql/expression/function/scalar/EsqlScalarFunction.java b/...main/java/org/elasticsearch/xpack/esql/expression/function/scalar/EsqlScalarFunction.java
@@ -7,10 +7,25 @@
 
 package org.elasticsearch.xpack.esql.expression.function.scalar;
 
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.function.scalar.ScalarFunction;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.evaluator.mapper.EvaluatorMapper;
+import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
+import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Greatest;
+import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Least;
+import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateDiff;
+import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateExtract;
+import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateFormat;
+import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateParse;
+import org.elasticsearch.xpack.esql.expression.function.scalar.date.DateTrunc;
+import org.elasticsearch.xpack.esql.expression.function.scalar.date.Now;
+import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
+import org.elasticsearch.xpack.esql.expression.function.scalar.string.Concat;
+import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToLower;
+import org.elasticsearch.xpack.esql.expression.function.scalar.string.ToUpper;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.InsensitiveEquals;
 
 import java.util.List;
 
@@ -25,6 +40,24 @@
  * </p>
  */
 public abstract class EsqlScalarFunction extends ScalarFunction implements EvaluatorMapper {
+    public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
+        return List.of(
+            Case.ENTRY,
+            Coalesce.ENTRY,
+            Concat.ENTRY,
+            Greatest.ENTRY,
+            InsensitiveEquals.ENTRY,
+            DateExtract.ENTRY,
+            DateDiff.ENTRY,
+            DateFormat.ENTRY,
+            DateParse.ENTRY,
+            DateTrunc.ENTRY,
+            Least.ENTRY,
+            Now.ENTRY,
+            ToLower.ENTRY,
+            ToUpper.ENTRY
+        );
+    }
 
     protected EsqlScalarFunction(Source source) {
         super(source);
@@ -38,5 +71,4 @@ protected EsqlScalarFunction(Source source, List<Expression> fields) {
     public Object fold() {
         return EvaluatorMapper.super.fold();
     }
-
 }
diff --git a/...ain/java/org/elasticsearch/xpack/esql/expression/function/scalar/UnaryScalarFunction.java b/...ain/java/org/elasticsearch/xpack/esql/expression/function/scalar/UnaryScalarFunction.java
@@ -12,6 +12,9 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
+import org.elasticsearch.xpack.esql.core.expression.predicate.logical.Not;
+import org.elasticsearch.xpack.esql.core.expression.predicate.nulls.IsNotNull;
+import org.elasticsearch.xpack.esql.core.expression.predicate.nulls.IsNull;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromBase64;
@@ -76,10 +79,13 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
             Cosh.ENTRY,
             Floor.ENTRY,
             FromBase64.ENTRY,
+            IsNotNull.ENTRY,
+            IsNull.ENTRY,
             Length.ENTRY,
             Log10.ENTRY,
             LTrim.ENTRY,
             Neg.ENTRY,
+            Not.ENTRY,
             RTrim.ENTRY,
             Signum.ENTRY,
             Sin.ENTRY,

diff --git a/...c/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java b/...c/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Case.java
@@ -7,6 +7,9 @@
 
 package org.elasticsearch.xpack.esql.expression.function.scalar.conditional;
 
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.ElementType;
@@ -27,8 +30,11 @@
 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.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput;
 import org.elasticsearch.xpack.esql.planner.PlannerUtils;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
@@ -37,8 +43,12 @@
 
 import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
 import static org.elasticsearch.xpack.esql.core.type.DataType.NULL;
+import static org.elasticsearch.xpack.esql.io.stream.PlanNameRegistry.PlanReader.readerFromPlanReader;
+import static org.elasticsearch.xpack.esql.io.stream.PlanNameRegistry.PlanWriter.writerFromPlanWriter;
 
 public final class Case extends EsqlScalarFunction {
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Case", Case::new);
+
     record Condition(Expression condition, Expression value) {}
 
     private final List<Condition> conditions;
@@ -110,6 +120,26 @@ public Case(
         elseValue = elseValueIsExplicit() ? children().get(children().size() - 1) : new Literal(source, null, NULL);
     }
 
+    private Case(StreamInput in) throws IOException {
+        this(
+            Source.readFrom((PlanStreamInput) in),
+            ((PlanStreamInput) in).readExpression(),
+            in.readCollectionAsList(readerFromPlanReader(PlanStreamInput::readExpression))
+        );
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        source().writeTo(out);
+        ((PlanStreamOutput) out).writeExpression(children().get(0));
+        out.writeCollection(children().subList(1, children().size()), writerFromPlanWriter(PlanStreamOutput::writeExpression));
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
     private boolean elseValueIsExplicit() {
         return children().size() % 2 == 1;
     }

diff --git a/...in/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Greatest.java b/...in/java/org/elasticsearch/xpack/esql/expression/function/scalar/conditional/Greatest.java
@@ -8,6 +8,9 @@
 package org.elasticsearch.xpack.esql.expression.function.scalar.conditional;
 
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.compute.ann.Evaluator;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
@@ -23,17 +26,24 @@
 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.multivalue.MvMax;
+import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.io.stream.PlanStreamOutput;
 
+import java.io.IOException;
 import java.util.List;
 import java.util.function.Function;
 import java.util.stream.Stream;
 
 import static org.elasticsearch.xpack.esql.core.type.DataType.NULL;
+import static org.elasticsearch.xpack.esql.io.stream.PlanNameRegistry.PlanReader.readerFromPlanReader;
+import static org.elasticsearch.xpack.esql.io.stream.PlanNameRegistry.PlanWriter.writerFromPlanWriter;
 
 /**
  * Returns the maximum value of multiple columns.
  */
 public class Greatest extends EsqlScalarFunction implements OptionalArgument {
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Greatest", Greatest::new);
+
     private DataType dataType;
 
     @FunctionInfo(
@@ -61,6 +71,26 @@ public Greatest(
         super(source, Stream.concat(Stream.of(first), rest.stream()).toList());
     }
 
+    private Greatest(StreamInput in) throws IOException {
+        this(
+            Source.readFrom((PlanStreamInput) in),
+            ((PlanStreamInput) in).readExpression(),
+            in.readCollectionAsList(readerFromPlanReader(PlanStreamInput::readExpression))
+        );
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        source().writeTo(out);
+        ((PlanStreamOutput) out).writeExpression(children().get(0));
+        out.writeCollection(children().subList(1, children().size()), writerFromPlanWriter(PlanStreamOutput::writeExpression));
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
     @Override
     public DataType dataType() {
         if (dataType == null) {