ESQL: Fix synthetic attribute pruning

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/plan/QueryPlan.java

@@ -35,6 +35,11 @@ public QueryPlan(Source source, List<PlanType> children) {
 
     public abstract List<Attribute> output();
 
+    /**
+     * The attributes required to be in the {@link QueryPlan#inputSet()} for this plan to be valid.
+     */
+    public abstract AttributeSet requiredInputSet();
+
     public AttributeSet outputSet() {
         if (lazyOutputSet == null) {
             lazyOutputSet = new AttributeSet(output());

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/io/stream/PlanNamedTypes.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.common.util.iterable.Iterables;
+import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.dissect.DissectParser;
 import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -168,7 +169,8 @@ static AggregateExec readAggregateExec(PlanStreamInput in) throws IOException {
             in.readPhysicalPlanNode(),
             in.readNamedWriteableCollectionAsList(Expression.class),
             in.readNamedWriteableCollectionAsList(NamedExpression.class),
-            in.readEnum(AggregateExec.Mode.class),
+            in.readEnum(AggregatorMode.class),
+            in.readNamedWriteableCollectionAsList(Attribute.class),
             in.readOptionalVInt()
         );
     }
@@ -179,6 +181,7 @@ static void writeAggregateExec(PlanStreamOutput out, AggregateExec aggregateExec
         out.writeNamedWriteableCollection(aggregateExec.groupings());
         out.writeNamedWriteableCollection(aggregateExec.aggregates());
         out.writeEnum(aggregateExec.getMode());
+        out.writeNamedWriteableCollection(aggregateExec.intermediateAttributes());
         out.writeOptionalVInt(aggregateExec.estimatedRowSize());
     }
 

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

@@ -468,6 +468,11 @@ 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 fail.
+                        // We 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. And this mapping needs to be kept
+                        // track of inside AggregateExec
+                        // Likely required for https://github.com/elastic/elasticsearch/issues/105436
                         tuple.v1().addAll(intermediateAttributes);
                         tuple.v2().add(stat);
                     }
@@ -520,6 +525,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

@@ -126,7 +126,7 @@ public static String temporaryName(Expression inner, Expression outer, int suffi
     }
 
     public static String locallyUniqueTemporaryName(String inner, String outer) {
-        return FieldAttribute.SYNTHETIC_ATTRIBUTE_NAME_PREFIX + inner + "$" + outer + "$" + new NameId();
+        return rawTemporaryName(inner, outer, (new NameId()).toString());
     }
 
     public static String rawTemporaryName(String inner, String outer, String suffix) {

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.optimizer;
 
+import org.elasticsearch.compute.aggregation.AggregatorMode;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
@@ -151,7 +152,7 @@ protected AttributeSet generates(PhysicalPlan physicalPlan) {
         @Override
         protected AttributeSet references(PhysicalPlan plan) {
             if (plan instanceof AggregateExec aggregate) {
-                if (aggregate.getMode() == AggregateExec.Mode.FINAL) {
+                if (aggregate.getMode() == AggregatorMode.FINAL) {
                     // lousy hack - need to generate the intermediate aggs yet the intermediateAggs method keep creating new IDs on each
                     // call
                     // in practice, the final aggregate should clearly declare the expected properties not hold on the original ones

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

@@ -11,12 +11,9 @@
 import org.elasticsearch.xpack.esql.common.Failure;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
-import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
-import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
-import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.rule.ParameterizedRuleExecutor;
 import org.elasticsearch.xpack.esql.core.rule.Rule;
 import org.elasticsearch.xpack.esql.core.rule.RuleExecutor;
@@ -25,15 +22,9 @@
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
-import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
-import org.elasticsearch.xpack.esql.plan.physical.EnrichExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
-import org.elasticsearch.xpack.esql.plan.physical.HashJoinExec;
-import org.elasticsearch.xpack.esql.plan.physical.MvExpandExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
-import org.elasticsearch.xpack.esql.plan.physical.ProjectExec;
-import org.elasticsearch.xpack.esql.plan.physical.RegexExtractExec;
 import org.elasticsearch.xpack.esql.plan.physical.UnaryExec;
 
 import java.util.ArrayList;
@@ -89,62 +80,26 @@ static class ProjectAwayColumns extends Rule<PhysicalPlan, PhysicalPlan> {
 
         @Override
         public PhysicalPlan apply(PhysicalPlan plan) {
-            var projectAll = new Holder<>(TRUE);
-            var keepCollecting = new Holder<>(TRUE);
-            var attributes = new AttributeSet();
-            var aliases = new AttributeMap<Expression>();
+            Holder<Boolean> keepTraversing = new Holder<>(TRUE);
+            // Invariant: if we add a projection with these attributes after the current plan node, the plan remains valid
+            // and the overall output will not change.
+            Holder<AttributeSet> requiredAttributes = new Holder<>(plan.outputSet());
 
-            return plan.transformDown(UnaryExec.class, p -> {
-                // no need for project all
-                if (p instanceof ProjectExec || p instanceof AggregateExec) {
-                    projectAll.set(FALSE);
+            // This will require updating should we choose to have non-unary execution plans in the future.
+            return plan.transformDown(UnaryExec.class, currentPlanNode -> {
+                if (keepTraversing.get() == false) {
+                    return currentPlanNode;
                 }
-                if (keepCollecting.get()) {
-                    p.forEachExpression(NamedExpression.class, ne -> {
-                        var attr = ne.toAttribute();
-                        // filter out attributes declared as aliases before
-                        if (ne instanceof Alias as) {
-                            aliases.put(attr, as.child());
-                            attributes.remove(attr);
-                        } else {
-                            // skip synthetically added attributes (the ones from AVG), see LogicalPlanOptimizer.SubstituteSurrogates
-                            if (attr.synthetic() == false && aliases.containsKey(attr) == false) {
-                                attributes.add(attr);
-                            }
-                        }
-                    });
-                    if (p instanceof RegexExtractExec ree) {
-                        attributes.removeAll(ree.extractedFields());
-                    }
-                    if (p instanceof MvExpandExec mvee) {
-                        attributes.remove(mvee.expanded());
-                    }
-                    if (p instanceof HashJoinExec join) {
-                        attributes.removeAll(join.addedFields());
-                        for (Attribute rhs : join.rightFields()) {
-                            if (join.leftFields().stream().anyMatch(x -> x.semanticEquals(rhs)) == false) {
-                                attributes.remove(rhs);
-                            }
-                        }
-                    }
-                    if (p instanceof EnrichExec ee) {
-                        for (NamedExpression enrichField : ee.enrichFields()) {
-                            // TODO: why is this different then the remove above?
-                            attributes.remove(enrichField instanceof Alias a ? a.child() : enrichField);
-                        }
-                    }
-                }
-                if (p instanceof ExchangeExec exec) {
-                    keepCollecting.set(FALSE);
+                if (currentPlanNode instanceof ExchangeExec exec) {
+                    keepTraversing.set(FALSE);
                     var child = exec.child();
                     // otherwise expect a Fragment
                     if (child instanceof FragmentExec fragmentExec) {
                         var logicalFragment = fragmentExec.fragment();
 
                         // no need for projection when dealing with aggs
                         if (logicalFragment instanceof Aggregate == false) {
-                            var selectAll = projectAll.get();
-                            var output = selectAll ? exec.child().output() : new ArrayList<>(attributes);
+                            List<Attribute> output = new ArrayList<>(requiredAttributes.get());
                             // if all the fields are filtered out, it's only the count that matters
                             // however until a proper fix (see https://github.com/elastic/elasticsearch/issues/98703)
                             // add a synthetic field (so it doesn't clash with the user defined one) to return a constant
@@ -156,19 +111,22 @@ public PhysicalPlan apply(PhysicalPlan plan) {
                                 output = Expressions.asAttributes(fields);
                             }
                             // add a logical projection (let the local replanning remove it if needed)
-                            p = exec.replaceChild(
-                                new FragmentExec(
-                                    Source.EMPTY,
-                                    new Project(logicalFragment.source(), logicalFragment, output),
-                                    fragmentExec.esFilter(),
-                                    fragmentExec.estimatedRowSize(),
-                                    fragmentExec.reducer()
-                                )
+                            FragmentExec newChild = new FragmentExec(
+                                Source.EMPTY,
+                                new Project(logicalFragment.source(), logicalFragment, output),
+                                fragmentExec.esFilter(),
+                                fragmentExec.estimatedRowSize(),
+                                fragmentExec.reducer()
                             );
+                            return new ExchangeExec(exec.source(), output, exec.isInBetweenAggs(), newChild);
                         }
                     }
+                } else {
+                    AttributeSet childOutput = currentPlanNode.inputSet();
+                    AttributeSet addedAttributes = currentPlanNode.outputSet().subtract(childOutput);
+                    requiredAttributes.set(requiredAttributes.get().subtract(addedAttributes).combine(currentPlanNode.requiredInputSet()));
                 }
-                return p;
+                return currentPlanNode;
             });
         }
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Aggregate.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
@@ -120,11 +121,24 @@ public boolean expressionsResolved() {
     @Override
     public List<Attribute> output() {
         if (lazyOutput == null) {
-            lazyOutput = mergeOutputAttributes(Expressions.asAttributes(aggregates()), emptyList());
+            lazyOutput = outputAttributes(aggregates);
         }
         return lazyOutput;
     }
 
+    public static List<Attribute> outputAttributes(List<? extends NamedExpression> aggregates) {
+        return mergeOutputAttributes(Expressions.asAttributes(aggregates), emptyList());
+    }
+
+    @Override
+    public AttributeSet requiredInputSet() {
+        return requiredInputAttributes(aggregates, groupings);
+    }
+
+    public static AttributeSet requiredInputAttributes(List<? extends NamedExpression> aggregates, List<? extends Expression> groupings) {
+        return Expressions.references(groupings).combine(Expressions.references(aggregates));
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(aggregateType, groupings, aggregates, child());

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Drop.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.plan.logical;
 
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -27,6 +28,11 @@ public List<NamedExpression> removals() {
         return removals;
     }
 
+    @Override
+    public AttributeSet requiredInputSet() {
+        return references();
+    }
+
     @Override
     public boolean expressionsResolved() {
         return Resolvables.resolved(removals);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Enrich.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.EmptyAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.NameId;
@@ -104,6 +105,11 @@ public Mode mode() {
         return mode;
     }
 
+    @Override
+    public AttributeSet requiredInputSet() {
+        return matchField.references();
+    }
+
     @Override
     public boolean expressionsResolved() {
         return policyName.resolved()

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

@@ -11,6 +11,8 @@
 import org.elasticsearch.xpack.esql.core.expression.Alias;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
+import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.NameId;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
@@ -47,6 +49,16 @@ public List<Attribute> output() {
         return lazyOutput;
     }
 
+    @Override
+    public AttributeSet requiredInputSet() {
+        return requiredAttributesFromChild(fields);
+    }
+
+    public static AttributeSet requiredAttributesFromChild(List<Alias> fields) {
+        AttributeSet generated = new AttributeSet(asAttributes(fields));
+        return Expressions.references(fields).subtract(generated);
+    }
+
     @Override
     public List<Attribute> generatedAttributes() {
         return asAttributes(fields);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Filter.java

@@ -6,6 +6,7 @@
  */
 package org.elasticsearch.xpack.esql.plan.logical;
 
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -26,6 +27,11 @@ public Filter(Source source, LogicalPlan child, Expression condition) {
         this.condition = condition;
     }
 
+    @Override
+    public AttributeSet requiredInputSet() {
+        return references();
+    }
+
     @Override
     protected NodeInfo<Filter> info() {
         return NodeInfo.create(this, Filter::new, child(), condition);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/InlineStats.java

@@ -88,6 +88,11 @@ public String getWriteableName() {
         return ENTRY.name;
     }
 
+    @Override
+    public AttributeSet requiredInputSet() {
+        return references();
+    }
+
     @Override
     protected NodeInfo<InlineStats> info() {
         return NodeInfo.create(this, InlineStats::new, child(), groupings, aggregates);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/LeafPlan.java

@@ -6,6 +6,7 @@
  */
 package org.elasticsearch.xpack.esql.plan.logical;
 
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 
 import java.util.Collections;
@@ -17,6 +18,11 @@ protected LeafPlan(Source source) {
         super(source, Collections.emptyList());
     }
 
+    @Override
+    public AttributeSet requiredInputSet() {
+        return AttributeSet.EMPTY;
+    }
+
     @Override
     public final LogicalPlan replaceChildren(List<LogicalPlan> newChildren) {
         throw new UnsupportedOperationException("this type of node doesn't have any children to replace");