ESQL: Fix synthetic attribute pruning (elastic#111413)

- Fix ProjectAwayColumns to handle synthetic attributes and mark synthetically introduced Aliases as synthetic again.
- Fix QueryPlan.references()
- Simplify ProjectAwayColumns.
- Simplify DependencyConsistency.
- Make AggregateExec track the intermediate attributes it actually outputs/requries.

Author: alex-spies

docs/changelog/111413.yaml

@@ -0,0 +1,6 @@
+pr: 111413
+summary: "ESQL: Fix synthetic attribute pruning"
+area: ES|QL
+type: bug
+issues:
+ - 105821

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -205,6 +205,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion GET_DATA_STREAMS_VERBOSE = def(8_735_00_0);
     public static final TransportVersion ESQL_ADD_INDEX_MODE_CONCRETE_INDICES = def(8_736_00_0);
     public static final TransportVersion UNASSIGNED_PRIMARY_COUNT_ON_CLUSTER_HEALTH = def(8_737_00_0);
+    public static final TransportVersion ESQL_AGGREGATE_EXEC_TRACKS_INTERMEDIATE_ATTRS = def(8_738_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -28,6 +28,11 @@
  * The rest are not as they are not part of the projection and thus are not part of the derived table.
  */
 public abstract class Attribute extends NamedExpression {
+    /**
+     * Changing this will break bwc with 8.15, see {@link FieldAttribute#fieldName()}.
+     */
+    protected static final String SYNTHETIC_ATTRIBUTE_NAME_PREFIX = "$$";
+
     public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         // TODO add UnsupportedAttribute when these are moved to the same project
         return List.of(FieldAttribute.ENTRY, MetadataAttribute.ENTRY, ReferenceAttribute.ENTRY);
@@ -49,6 +54,10 @@ public Attribute(Source source, String name, Nullability nullability, NameId id,
         this.nullability = nullability;
     }
 
+    public static String rawTemporaryName(String inner, String outer, String suffix) {
+        return SYNTHETIC_ATTRIBUTE_NAME_PREFIX + inner + "$" + outer + "$" + suffix;
+    }
+
     @Override
     public final Expression replaceChildren(List<Expression> newChildren) {
         throw new UnsupportedOperationException("this type of node doesn't have any children to replace");
@@ -123,7 +132,7 @@ public boolean equals(Object obj) {
 
     @Override
     public String toString() {
-        return name() + "{" + label() + "}" + "#" + id();
+        return name() + "{" + label() + (synthetic() ? "$" : "") + "}" + "#" + id();
     }
 
     @Override

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

@@ -30,9 +30,6 @@
  * - nestedParent - if nested, what's the parent (which might not be the immediate one)
  */
 public class FieldAttribute extends TypedAttribute {
-    // TODO: This constant should not be used if possible; use .synthetic()
-    // https://github.com/elastic/elasticsearch/issues/105821
-    public static final String SYNTHETIC_ATTRIBUTE_NAME_PREFIX = "$$";
 
     static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         Attribute.class,
@@ -52,6 +49,10 @@ public FieldAttribute(Source source, FieldAttribute parent, String name, EsField
         this(source, parent, name, field, Nullability.TRUE, null, false);
     }
 
+    public FieldAttribute(Source source, FieldAttribute parent, String name, EsField field, boolean synthetic) {
+        this(source, parent, name, field, Nullability.TRUE, null, synthetic);
+    }
+
     public FieldAttribute(
         Source source,
         FieldAttribute parent,
@@ -64,7 +65,11 @@ public FieldAttribute(
         this(source, parent, name, field.getDataType(), field, nullability, id, synthetic);
     }
 
-    public FieldAttribute(
+    /**
+     * Used only for testing. Do not use this otherwise, as an explicitly set type will be ignored the next time this FieldAttribute is
+     * {@link FieldAttribute#clone}d.
+     */
+    FieldAttribute(
         Source source,
         FieldAttribute parent,
         String name,
@@ -147,28 +152,7 @@ public String getWriteableName() {
 
     @Override
     protected NodeInfo<FieldAttribute> info() {
-        return NodeInfo.create(
-            this,
-            (source, parent1, name, type, field1, qualifier, nullability, id, synthetic) -> new FieldAttribute(
-                source,
-                parent1,
-                name,
-                type,
-                field1,
-                qualifier,
-                nullability,
-                id,
-                synthetic
-            ),
-            parent,
-            name(),
-            dataType(),
-            field,
-            (String) null,
-            nullable(),
-            id(),
-            synthetic()
-        );
+        return NodeInfo.create(this, FieldAttribute::new, parent, name(), dataType(), field, (String) null, nullable(), id(), synthetic());
     }
 
     public FieldAttribute parent() {
@@ -185,9 +169,9 @@ public String path() {
     public String fieldName() {
         // Before 8.15, the field name was the same as the attribute's name.
         // On later versions, the attribute can be renamed when creating synthetic attributes.
-        // TODO: We should use synthetic() to check for that case.
-        // https://github.com/elastic/elasticsearch/issues/105821
-        if (name().startsWith(SYNTHETIC_ATTRIBUTE_NAME_PREFIX) == false) {
+        // Because until 8.15, we couldn't set `synthetic` to true due to a bug, in that version such FieldAttributes are marked by their
+        // name starting with `$$`.
+        if ((synthetic() || name().startsWith(SYNTHETIC_ATTRIBUTE_NAME_PREFIX)) == false) {
             return name();
         }
         return Strings.hasText(path) ? path + "." + field.getName() : field.getName();
@@ -211,9 +195,15 @@ private FieldAttribute innerField(EsField type) {
 
     @Override
     protected Attribute clone(Source source, String name, DataType type, Nullability nullability, NameId id, boolean synthetic) {
+        // Ignore `type`, this must be the same as the field's type.
         return new FieldAttribute(source, parent, name, field, nullability, id, synthetic);
     }
 
+    @Override
+    public Attribute withDataType(DataType type) {
+        throw new UnsupportedOperationException("FieldAttribute obtains its type from the contained EsField.");
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(super.hashCode(), path, field);

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

@@ -51,6 +51,10 @@ public NameId id() {
         return id;
     }
 
+    /**
+     * Synthetic named expressions are not user defined and usually created during optimizations and substitutions, e.g. when turning
+     * {@code ... | STATS x = avg(2*field)} into {@code ... | EVAL $$synth$attribute = 2*field | STATS x = avg($$synth$attribute)}.
+     */
     public boolean synthetic() {
         return synthetic;
     }

x-pack/plugin/esql-core/src/test/java/org/elasticsearch/xpack/esql/core/expression/FieldAttributeTestUtils.java

@@ -0,0 +1,27 @@
+/*
+ * 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.core.expression;
+
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.core.type.EsField;
+
+public class FieldAttributeTestUtils {
+    public static final FieldAttribute newFieldAttributeWithType(
+        Source source,
+        FieldAttribute parent,
+        String name,
+        DataType type,
+        EsField field,
+        Nullability nullability,
+        NameId id,
+        boolean synthetic
+    ) {
+        return new FieldAttribute(source, parent, name, type, field, nullability, id, synthetic);
+    }
+}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/aggregation/AggregatorMode.java

@@ -9,13 +9,24 @@
 
 public enum AggregatorMode {
 
+    /**
+     * Maps raw inputs to intermediate outputs.
+     */
     INITIAL(false, true),
 
+    /**
+     * Maps intermediate inputs to intermediate outputs.
+     */
     INTERMEDIATE(true, true),
 
+    /**
+     * Maps intermediate inputs to final outputs.
+     */
     FINAL(true, false),
 
-    // most useful for testing
+    /**
+     * Maps raw inputs to final outputs. Most useful for testing.
+     */
     SINGLE(false, false);
 
     private final boolean inputPartial;

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

@@ -26,7 +26,6 @@
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
-import org.elasticsearch.xpack.esql.core.expression.NameId;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
@@ -65,7 +64,6 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.EsqlArithmeticOperation;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.In;
 import org.elasticsearch.xpack.esql.index.EsIndex;
-import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.plan.TableIdentifier;
 import org.elasticsearch.xpack.esql.plan.logical.Drop;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -1249,12 +1247,10 @@ private Expression createIfDoesNotAlreadyExist(
             List<FieldAttribute> unionFieldAttributes
         ) {
             // Generate new ID for the field and suffix it with the data type to maintain unique attribute names.
-            String unionTypedFieldName = LogicalPlanOptimizer.rawTemporaryName(
-                fa.name(),
-                "converted_to",
-                resolvedField.getDataType().typeName()
-            );
-            FieldAttribute unionFieldAttribute = new FieldAttribute(fa.source(), fa.parent(), unionTypedFieldName, resolvedField);
+            // NOTE: The name has to start with $$ to not break bwc with 8.15 - in that version, this is how we had to mark this as
+            // synthetic to work around a bug.
+            String unionTypedFieldName = Attribute.rawTemporaryName(fa.name(), "converted_to", resolvedField.getDataType().typeName());
+            FieldAttribute unionFieldAttribute = new FieldAttribute(fa.source(), fa.parent(), unionTypedFieldName, resolvedField, true);
             int existingIndex = unionFieldAttributes.indexOf(unionFieldAttribute);
             if (existingIndex >= 0) {
                 // Do not generate multiple name/type combinations with different IDs
@@ -1280,8 +1276,16 @@ private MultiTypeEsField resolvedMultiTypeEsField(FieldAttribute fa, HashMap<Typ
 
         private Expression typeSpecificConvert(AbstractConvertFunction convert, Source source, DataType type, InvalidMappedField mtf) {
             EsField field = new EsField(mtf.getName(), type, mtf.getProperties(), mtf.isAggregatable());
-            NameId id = ((FieldAttribute) convert.field()).id();
-            FieldAttribute resolvedAttr = new FieldAttribute(source, null, field.getName(), field, Nullability.TRUE, id, false);
+            FieldAttribute originalFieldAttr = (FieldAttribute) convert.field();
+            FieldAttribute resolvedAttr = new FieldAttribute(
+                source,
+                originalFieldAttr.parent(),
+                originalFieldAttr.name(),
+                field,
+                originalFieldAttr.nullable(),
+                originalFieldAttr.id(),
+                true
+            );
             return convert.replaceChildren(Collections.singletonList(resolvedAttr));
         }
     }
@@ -1330,11 +1334,11 @@ private static LogicalPlan planWithoutSyntheticAttributes(LogicalPlan plan) {
             List<Attribute> newOutput = new ArrayList<>(output.size());
 
             for (Attribute attr : output) {
-                // TODO: this should really use .synthetic()
-                // https://github.com/elastic/elasticsearch/issues/105821
-                if (attr.name().startsWith(FieldAttribute.SYNTHETIC_ATTRIBUTE_NAME_PREFIX) == false) {
-                    newOutput.add(attr);
+                // Do not let the synthetic union type field attributes end up in the final output.
+                if (attr.synthetic() && attr instanceof FieldAttribute) {
+                    continue;
                 }
+                newOutput.add(attr);
             }
 
             return newOutput.size() == output.size() ? plan : new Project(Source.EMPTY, plan, newOutput);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java

@@ -552,6 +552,10 @@ private Tuple<List<Attribute>, List<Stat>> pushableStats(AggregateExec aggregate
                             singletonList(agg),
                             emptyList()
                         );
+                        // TODO: the attributes have been recreated here; they will have wrong name ids, and the dependency check will
+                        // probably fail when we fix https://github.com/elastic/elasticsearch/issues/105436.
+                        // We may need to refactor AbstractPhysicalOperationProviders.intermediateAttributes so it doesn't return just
+                        // a list of attributes, but a mapping from the logical to the physical attributes.
                         tuple.v1().addAll(intermediateAttributes);
                         tuple.v2().add(stat);
                     }
@@ -604,6 +608,7 @@ && allowedForDocValues(fieldAttribute, agg, foundAttributes)) {
                             agg.groupings(),
                             orderedAggregates,
                             agg.getMode(),
+                            agg.intermediateAttributes(),
                             agg.estimatedRowSize()
                         );
                     }

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

@@ -15,7 +15,6 @@
 import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
-import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.NameId;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
@@ -123,15 +122,11 @@ public LogicalPlanOptimizer(LogicalOptimizerContext optimizerContext) {
     public static String temporaryName(Expression inner, Expression outer, int suffix) {
         String in = toString(inner);
         String out = toString(outer);
-        return rawTemporaryName(in, out, String.valueOf(suffix));
+        return Attribute.rawTemporaryName(in, out, String.valueOf(suffix));
     }
 
     public static String locallyUniqueTemporaryName(String inner, String outer) {
-        return FieldAttribute.SYNTHETIC_ATTRIBUTE_NAME_PREFIX + inner + "$" + outer + "$" + new NameId();
-    }
-
-    public static String rawTemporaryName(String inner, String outer, String suffix) {
-        return FieldAttribute.SYNTHETIC_ATTRIBUTE_NAME_PREFIX + inner + "$" + outer + "$" + suffix;
+        return Attribute.rawTemporaryName(inner, outer, (new NameId()).toString());
     }
 
     static String toString(Expression ex) {
@@ -373,9 +368,7 @@ private static AttributeReplacement renameAttributesInExpressions(
                 if (attributeNamesToRename.contains(attr.name())) {
                     Alias renamedAttribute = aliasesForReplacedAttributes.computeIfAbsent(attr, a -> {
                         String tempName = locallyUniqueTemporaryName(a.name(), "temp_name");
-                        // TODO: this should be synthetic
-                        // blocked on https://github.com/elastic/elasticsearch/issues/98703
-                        return new Alias(a.source(), tempName, a, null, false);
+                        return new Alias(a.source(), tempName, a, null, true);
                     });
                     return renamedAttribute.toAttribute();
                 }