Browse Source

ESQL: Use Alias instead of NamedExpression on Eval and Row (#98859)

Field properties in Eval and Row are always a list of Aliases yet this
 is obfuscated by exposing them as a list of NamedExpressions which
 makes consumption difficult (trying to cast it to an Alias else throw
 an exception).
Costin Leau 2 years ago
parent
commit
481b6a3369

+ 3 - 6
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -348,13 +348,10 @@ public class Analyzer extends ParameterizedRuleExecutor<LogicalPlan, AnalyzerCon
 
 
         private LogicalPlan resolveEval(Eval eval, List<Attribute> childOutput) {
         private LogicalPlan resolveEval(Eval eval, List<Attribute> childOutput) {
             List<Attribute> allResolvedInputs = new ArrayList<>(childOutput);
             List<Attribute> allResolvedInputs = new ArrayList<>(childOutput);
-            List<NamedExpression> newFields = new ArrayList<>();
+            List<Alias> newFields = new ArrayList<>();
             boolean changed = false;
             boolean changed = false;
-            for (NamedExpression field : eval.fields()) {
-                NamedExpression result = (NamedExpression) field.transformUp(
-                    UnresolvedAttribute.class,
-                    ua -> resolveAttribute(ua, allResolvedInputs)
-                );
+            for (Alias field : eval.fields()) {
+                Alias result = (Alias) field.transformUp(UnresolvedAttribute.class, ua -> resolveAttribute(ua, allResolvedInputs));
 
 
                 changed |= result != field;
                 changed |= result != field;
                 newFields.add(result);
                 newFields.add(result);

+ 3 - 3
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java

@@ -228,9 +228,9 @@ public class Verifier {
 
 
     private static Collection<Failure> validateRow(Row row) {
     private static Collection<Failure> validateRow(Row row) {
         List<Failure> failures = new ArrayList<>(row.fields().size());
         List<Failure> failures = new ArrayList<>(row.fields().size());
-        row.fields().forEach(o -> {
-            if (EsqlDataTypes.isRepresentable(o.dataType()) == false && o instanceof Alias a) {
-                failures.add(fail(o, "cannot use [{}] directly in a row assignment", a.child().sourceText()));
+        row.fields().forEach(a -> {
+            if (EsqlDataTypes.isRepresentable(a.dataType()) == false) {
+                failures.add(fail(a, "cannot use [{}] directly in a row assignment", a.child().sourceText()));
             }
             }
         });
         });
         return failures;
         return failures;

+ 14 - 6
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java

@@ -438,12 +438,12 @@ public final class PlanNamedTypes {
     }
     }
 
 
     static EvalExec readEvalExec(PlanStreamInput in) throws IOException {
     static EvalExec readEvalExec(PlanStreamInput in) throws IOException {
-        return new EvalExec(Source.EMPTY, in.readPhysicalPlanNode(), readNamedExpressions(in));
+        return new EvalExec(Source.EMPTY, in.readPhysicalPlanNode(), readAliases(in));
     }
     }
 
 
     static void writeEvalExec(PlanStreamOutput out, EvalExec evalExec) throws IOException {
     static void writeEvalExec(PlanStreamOutput out, EvalExec evalExec) throws IOException {
         out.writePhysicalPlanNode(evalExec.child());
         out.writePhysicalPlanNode(evalExec.child());
-        writeNamedExpressions(out, evalExec.fields());
+        writeAliases(out, evalExec.fields());
     }
     }
 
 
     static EnrichExec readEnrichExec(PlanStreamInput in) throws IOException {
     static EnrichExec readEnrichExec(PlanStreamInput in) throws IOException {
@@ -582,12 +582,12 @@ public final class PlanNamedTypes {
     }
     }
 
 
     static RowExec readRowExec(PlanStreamInput in) throws IOException {
     static RowExec readRowExec(PlanStreamInput in) throws IOException {
-        return new RowExec(Source.EMPTY, readNamedExpressions(in));
+        return new RowExec(Source.EMPTY, readAliases(in));
     }
     }
 
 
     static void writeRowExec(PlanStreamOutput out, RowExec rowExec) throws IOException {
     static void writeRowExec(PlanStreamOutput out, RowExec rowExec) throws IOException {
         assert rowExec.children().size() == 0;
         assert rowExec.children().size() == 0;
-        writeNamedExpressions(out, rowExec.fields());
+        writeAliases(out, rowExec.fields());
     }
     }
 
 
     @SuppressWarnings("unchecked")
     @SuppressWarnings("unchecked")
@@ -655,12 +655,12 @@ public final class PlanNamedTypes {
     }
     }
 
 
     static Eval readEval(PlanStreamInput in) throws IOException {
     static Eval readEval(PlanStreamInput in) throws IOException {
-        return new Eval(Source.EMPTY, in.readLogicalPlanNode(), readNamedExpressions(in));
+        return new Eval(Source.EMPTY, in.readLogicalPlanNode(), readAliases(in));
     }
     }
 
 
     static void writeEval(PlanStreamOutput out, Eval eval) throws IOException {
     static void writeEval(PlanStreamOutput out, Eval eval) throws IOException {
         out.writeLogicalPlanNode(eval.child());
         out.writeLogicalPlanNode(eval.child());
-        writeNamedExpressions(out, eval.fields());
+        writeAliases(out, eval.fields());
     }
     }
 
 
     static Enrich readEnrich(PlanStreamInput in) throws IOException {
     static Enrich readEnrich(PlanStreamInput in) throws IOException {
@@ -772,6 +772,14 @@ public final class PlanNamedTypes {
         out.writeCollection(namedExpressions, writerFromPlanWriter(PlanStreamOutput::writeNamedExpression));
         out.writeCollection(namedExpressions, writerFromPlanWriter(PlanStreamOutput::writeNamedExpression));
     }
     }
 
 
+    private static List<Alias> readAliases(PlanStreamInput in) throws IOException {
+        return in.readList(readerFromPlanReader(PlanNamedTypes::readAlias));
+    }
+
+    static void writeAliases(PlanStreamOutput out, List<Alias> aliases) throws IOException {
+        out.writeCollection(aliases, writerFromPlanWriter(PlanNamedTypes::writeAlias));
+    }
+
     static FieldAttribute readFieldAttribute(PlanStreamInput in) throws IOException {
     static FieldAttribute readFieldAttribute(PlanStreamInput in) throws IOException {
         return new FieldAttribute(
         return new FieldAttribute(
             Source.EMPTY,
             Source.EMPTY,

+ 1 - 1
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalLogicalPlanOptimizer.java

@@ -88,7 +88,7 @@ public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<Logical
             else if (plan instanceof Project project) {
             else if (plan instanceof Project project) {
                 var projections = project.projections();
                 var projections = project.projections();
                 List<NamedExpression> newProjections = new ArrayList<>(projections.size());
                 List<NamedExpression> newProjections = new ArrayList<>(projections.size());
-                List<NamedExpression> literals = new ArrayList<>();
+                List<Alias> literals = new ArrayList<>();
 
 
                 for (NamedExpression projection : projections) {
                 for (NamedExpression projection : projections) {
                     if (projection instanceof FieldAttribute f && stats.exists(f.qualifiedName()) == false) {
                     if (projection instanceof FieldAttribute f && stats.exists(f.qualifiedName()) == false) {

+ 1 - 1
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java

@@ -141,7 +141,7 @@ public class LogicalPlanOptimizer extends RuleExecutor<LogicalPlan> {
             // existing aggregate and their respective attributes
             // existing aggregate and their respective attributes
             Map<AggregateFunction, Attribute> aggFuncToAttr = new HashMap<>();
             Map<AggregateFunction, Attribute> aggFuncToAttr = new HashMap<>();
             // surrogate functions eval
             // surrogate functions eval
-            List<NamedExpression> transientEval = new ArrayList<>();
+            List<Alias> transientEval = new ArrayList<>();
             boolean changed = false;
             boolean changed = false;
 
 
             // first pass to check existing aggregates (to avoid duplication and alias waste)
             // first pass to check existing aggregates (to avoid duplication and alias waste)

+ 4 - 4
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java

@@ -194,7 +194,7 @@ public class LogicalPlanBuilder extends ExpressionBuilder {
 
 
     @Override
     @Override
     public PlanFactory visitStatsCommand(EsqlBaseParser.StatsCommandContext ctx) {
     public PlanFactory visitStatsCommand(EsqlBaseParser.StatsCommandContext ctx) {
-        List<NamedExpression> aggregates = visitFields(ctx.fields());
+        List<NamedExpression> aggregates = new ArrayList<>(visitFields(ctx.fields()));
         List<NamedExpression> groupings = visitGrouping(ctx.grouping());
         List<NamedExpression> groupings = visitGrouping(ctx.grouping());
         if (aggregates.isEmpty() && groupings.isEmpty()) {
         if (aggregates.isEmpty() && groupings.isEmpty()) {
             throw new ParsingException(source(ctx), "At least one aggregation or grouping expression required in [{}]", ctx.getText());
             throw new ParsingException(source(ctx), "At least one aggregation or grouping expression required in [{}]", ctx.getText());
@@ -215,7 +215,7 @@ public class LogicalPlanBuilder extends ExpressionBuilder {
 
 
     @Override
     @Override
     public PlanFactory visitInlinestatsCommand(EsqlBaseParser.InlinestatsCommandContext ctx) {
     public PlanFactory visitInlinestatsCommand(EsqlBaseParser.InlinestatsCommandContext ctx) {
-        List<NamedExpression> aggregates = visitFields(ctx.fields());
+        List<NamedExpression> aggregates = new ArrayList<>(visitFields(ctx.fields()));
         List<NamedExpression> groupings = visitGrouping(ctx.grouping());
         List<NamedExpression> groupings = visitGrouping(ctx.grouping());
         aggregates.addAll(groupings);
         aggregates.addAll(groupings);
         return input -> new InlineStats(source(ctx), input, new ArrayList<>(groupings), aggregates);
         return input -> new InlineStats(source(ctx), input, new ArrayList<>(groupings), aggregates);
@@ -228,8 +228,8 @@ public class LogicalPlanBuilder extends ExpressionBuilder {
     }
     }
 
 
     @Override
     @Override
-    public List<NamedExpression> visitFields(EsqlBaseParser.FieldsContext ctx) {
-        return ctx != null ? visitList(this, ctx.field(), NamedExpression.class) : new ArrayList<>();
+    public List<Alias> visitFields(EsqlBaseParser.FieldsContext ctx) {
+        return ctx != null ? visitList(this, ctx.field(), Alias.class) : new ArrayList<>();
     }
     }
 
 
     @Override
     @Override

+ 4 - 4
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Eval.java

@@ -8,8 +8,8 @@
 package org.elasticsearch.xpack.esql.plan.logical;
 package org.elasticsearch.xpack.esql.plan.logical;
 
 
 import org.elasticsearch.xpack.ql.capabilities.Resolvables;
 import org.elasticsearch.xpack.ql.capabilities.Resolvables;
+import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.Attribute;
-import org.elasticsearch.xpack.ql.expression.NamedExpression;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.ql.plan.logical.UnaryPlan;
 import org.elasticsearch.xpack.ql.plan.logical.UnaryPlan;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
@@ -22,14 +22,14 @@ import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutp
 
 
 public class Eval extends UnaryPlan {
 public class Eval extends UnaryPlan {
 
 
-    private final List<NamedExpression> fields;
+    private final List<Alias> fields;
 
 
-    public Eval(Source source, LogicalPlan child, List<NamedExpression> fields) {
+    public Eval(Source source, LogicalPlan child, List<Alias> fields) {
         super(source, child);
         super(source, child);
         this.fields = fields;
         this.fields = fields;
     }
     }
 
 
-    public List<NamedExpression> fields() {
+    public List<Alias> fields() {
         return fields;
         return fields;
     }
     }
 
 

+ 4 - 4
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Row.java

@@ -8,9 +8,9 @@
 package org.elasticsearch.xpack.esql.plan.logical;
 package org.elasticsearch.xpack.esql.plan.logical;
 
 
 import org.elasticsearch.xpack.ql.capabilities.Resolvables;
 import org.elasticsearch.xpack.ql.capabilities.Resolvables;
+import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.Expressions;
-import org.elasticsearch.xpack.ql.expression.NamedExpression;
 import org.elasticsearch.xpack.ql.plan.logical.LeafPlan;
 import org.elasticsearch.xpack.ql.plan.logical.LeafPlan;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
@@ -21,14 +21,14 @@ import java.util.Objects;
 
 
 public class Row extends LeafPlan {
 public class Row extends LeafPlan {
 
 
-    private final List<NamedExpression> fields;
+    private final List<Alias> fields;
 
 
-    public Row(Source source, List<NamedExpression> fields) {
+    public Row(Source source, List<Alias> fields) {
         super(source);
         super(source);
         this.fields = fields;
         this.fields = fields;
     }
     }
 
 
-    public List<NamedExpression> fields() {
+    public List<Alias> fields() {
         return fields;
         return fields;
     }
     }
 
 

+ 4 - 4
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/EvalExec.java

@@ -7,8 +7,8 @@
 
 
 package org.elasticsearch.xpack.esql.plan.physical;
 package org.elasticsearch.xpack.esql.plan.physical;
 
 
+import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.Attribute;
-import org.elasticsearch.xpack.ql.expression.NamedExpression;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
 import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.tree.Source;
 
 
@@ -18,14 +18,14 @@ import java.util.Objects;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
 
 public class EvalExec extends UnaryExec implements EstimatesRowSize {
 public class EvalExec extends UnaryExec implements EstimatesRowSize {
-    private final List<NamedExpression> fields;
+    private final List<Alias> fields;
 
 
-    public EvalExec(Source source, PhysicalPlan child, List<NamedExpression> fields) {
+    public EvalExec(Source source, PhysicalPlan child, List<Alias> fields) {
         super(source, child);
         super(source, child);
         this.fields = fields;
         this.fields = fields;
     }
     }
 
 
-    public List<NamedExpression> fields() {
+    public List<Alias> fields() {
         return fields;
         return fields;
     }
     }
 
 

+ 4 - 4
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/RowExec.java

@@ -7,9 +7,9 @@
 
 
 package org.elasticsearch.xpack.esql.plan.physical;
 package org.elasticsearch.xpack.esql.plan.physical;
 
 
+import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.Expressions;
 import org.elasticsearch.xpack.ql.expression.Expressions;
-import org.elasticsearch.xpack.ql.expression.NamedExpression;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
 import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.tree.Source;
 
 
@@ -17,14 +17,14 @@ import java.util.List;
 import java.util.Objects;
 import java.util.Objects;
 
 
 public class RowExec extends LeafExec {
 public class RowExec extends LeafExec {
-    private final List<NamedExpression> fields;
+    private final List<Alias> fields;
 
 
-    public RowExec(Source source, List<NamedExpression> fields) {
+    public RowExec(Source source, List<Alias> fields) {
         super(source);
         super(source);
         this.fields = fields;
         this.fields = fields;
     }
     }
 
 
-    public List<NamedExpression> fields() {
+    public List<Alias> fields() {
         return fields;
         return fields;
     }
     }
 
 

+ 4 - 14
x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java

@@ -383,15 +383,11 @@ public class LocalExecutionPlanner {
     private PhysicalOperation planEval(EvalExec eval, LocalExecutionPlannerContext context) {
     private PhysicalOperation planEval(EvalExec eval, LocalExecutionPlannerContext context) {
         PhysicalOperation source = plan(eval.child(), context);
         PhysicalOperation source = plan(eval.child(), context);
 
 
-        for (NamedExpression namedExpression : eval.fields()) {
+        for (Alias field : eval.fields()) {
             Supplier<ExpressionEvaluator> evaluatorSupplier;
             Supplier<ExpressionEvaluator> evaluatorSupplier;
-            if (namedExpression instanceof Alias alias) {
-                evaluatorSupplier = EvalMapper.toEvaluator(alias.child(), source.layout);
-            } else {
-                throw new EsqlIllegalArgumentException("source fields for eval nodes have to be aliases");
-            }
+            evaluatorSupplier = EvalMapper.toEvaluator(field.child(), source.layout);
             Layout.Builder layout = source.layout.builder();
             Layout.Builder layout = source.layout.builder();
-            layout.appendChannel(namedExpression.toAttribute().id());
+            layout.appendChannel(field.toAttribute().id());
             source = source.with(new EvalOperatorFactory(evaluatorSupplier), layout.build());
             source = source.with(new EvalOperatorFactory(evaluatorSupplier), layout.build());
         }
         }
         return source;
         return source;
@@ -483,13 +479,7 @@ public class LocalExecutionPlanner {
     }
     }
 
 
     private PhysicalOperation planRow(RowExec row, LocalExecutionPlannerContext context) {
     private PhysicalOperation planRow(RowExec row, LocalExecutionPlannerContext context) {
-        List<Object> obj = row.fields().stream().map(f -> {
-            if (f instanceof Alias) {
-                return ((Alias) f).child().fold();
-            } else {
-                return f.fold();
-            }
-        }).toList();
+        List<Object> obj = row.fields().stream().map(f -> f.child().fold()).toList();
         Layout.Builder layout = new Layout.Builder();
         Layout.Builder layout = new Layout.Builder();
         var output = row.output();
         var output = row.output();
         for (Attribute attribute : output) {
         for (Attribute attribute : output) {

+ 1 - 1
x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java

@@ -117,7 +117,7 @@ public class AnalyzerTests extends ESTestCase {
         var eval = as(limit.child(), Eval.class);
         var eval = as(limit.child(), Eval.class);
 
 
         assertEquals(1, eval.fields().size());
         assertEquals(1, eval.fields().size());
-        Alias eeField = (Alias) eval.fields().get(0);
+        Alias eeField = eval.fields().get(0);
         assertEquals("ee", eeField.name());
         assertEquals("ee", eeField.name());
         assertEquals("e", ((ReferenceAttribute) eeField.child()).name());
         assertEquals("e", ((ReferenceAttribute) eeField.child()).name());
 
 

+ 1 - 1
x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypesTests.java

@@ -149,7 +149,7 @@ public class PlanNamedTypesTests extends ESTestCase {
         BytesStreamOutput bso = new BytesStreamOutput();
         BytesStreamOutput bso = new BytesStreamOutput();
         bso.writeString("hello");
         bso.writeString("hello");
         PlanStreamOutput out = new PlanStreamOutput(bso, planNameRegistry);
         PlanStreamOutput out = new PlanStreamOutput(bso, planNameRegistry);
-        var plan = new RowExec(Source.EMPTY, List.of(field("foo", DataTypes.LONG)));
+        var plan = new RowExec(Source.EMPTY, List.of(new Alias(Source.EMPTY, "foo", field("field", DataTypes.LONG))));
         out.writePhysicalPlanNode(plan);
         out.writePhysicalPlanNode(plan);
         bso.writeVInt(11_345);
         bso.writeVInt(11_345);