Adding support for pruning columns in fork branches

Author: pmpailis

docs/changelog/137907.yaml

@@ -0,0 +1,6 @@
+pr: 137907
+summary: Prune columns when using fork
+area: ES|QL
+type: bug
+issues:
+ - 136365

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

@@ -212,11 +212,15 @@ public static Builder builder(int expectedSize) {
         return new Builder(AttributeMap.builder(expectedSize));
     }
 
+    public static Builder forkBuilder() {
+        return new ForkBuilder(AttributeMap.builder());
+    }
+
     public static class Builder {
 
-        private final AttributeMap.Builder<Object> mapBuilder;
+        protected final AttributeMap.Builder<Object> mapBuilder;
 
-        private Builder(AttributeMap.Builder<Object> mapBuilder) {
+        protected Builder(AttributeMap.Builder<Object> mapBuilder) {
             this.mapBuilder = mapBuilder;
         }
 
@@ -266,4 +270,25 @@ public void clear() {
             mapBuilder.keySet().clear();
         }
     }
+
+    /**
+     * This class extends {@code Builder}, but its {@code contains} method also matches {@code NamedExpression} instances by their name.
+     * This is useful for Fork plans, where branches may have different Attribute IDs but share a common output schema,
+     * allowing equality checks of used attributes based on their names.
+     */
+    public static class ForkBuilder extends Builder {
+
+        private ForkBuilder(AttributeMap.Builder<Object> mapBuilder) {
+            super(mapBuilder);
+        }
+
+        @Override
+        public boolean contains(Object o) {
+            if (super.contains(o)) {
+                return true;
+            }
+            return o instanceof NamedExpression
+                && mapBuilder.keySet().stream().anyMatch(x -> x.name().equals(((NamedExpression) o).name()));
+        }
+    }
 }

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

@@ -34,6 +34,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateNullable;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropgateUnmappedFields;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumnsInForkBranches;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneEmptyAggregates;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneFilters;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneLiteralsInOrderBy;
@@ -215,6 +216,7 @@ protected static Batch<LogicalPlan> operators() {
             new PruneFilters(),
             new PruneColumns(),
             new PruneLiteralsInOrderBy(),
+            new PruneColumnsInForkBranches(),
             new PushDownAndCombineLimits(),
             new PushLimitToKnn(),
             new PushDownAndCombineFilters(),

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java

@@ -25,14 +25,17 @@
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.Sample;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
+import org.elasticsearch.xpack.esql.plan.logical.UnionAll;
 import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
 import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
 import org.elasticsearch.xpack.esql.plan.logical.local.LocalSupplier;
 import org.elasticsearch.xpack.esql.planner.PlannerUtils;
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 import static org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneEmptyPlans.skipPlan;
 
@@ -46,9 +49,8 @@ public LogicalPlan apply(LogicalPlan plan) {
         return pruneColumns(plan, plan.outputSet().asBuilder(), false);
     }
 
-    private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder used, boolean inlineJoin) {
+    static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder used, boolean inlineJoin) {
         Holder<Boolean> forkPresent = new Holder<>(false);
-
         // while going top-to-bottom (upstream)
         return plan.transformDown(p -> {
             // Note: It is NOT required to do anything special for binary plans like JOINs, except INLINE STATS. It is perfectly fine that
@@ -58,17 +60,13 @@ private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder u
             // same index fields will have different name ids in the left and right hand sides - as in the extreme example
             // `FROM lookup_idx | LOOKUP JOIN lookup_idx ON key_field`.
 
-            // TODO: revisit with every new command
-            // skip nodes that simply pass the input through and use no references
-            if (p instanceof Limit || p instanceof Sample) {
+            if (forkPresent.get()) {
                 return p;
             }
 
-            if (p instanceof Fork) {
-                forkPresent.set(true);
-            }
-            // pruning columns for Fork branches can have the side effect of having misaligned outputs
-            if (forkPresent.get()) {
+            // TODO: revisit with every new command
+            // skip nodes that simply pass the input through and use no references
+            if (p instanceof Limit || p instanceof Sample) {
                 return p;
             }
 
@@ -83,6 +81,10 @@ private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder u
                     case Eval eval -> pruneColumnsInEval(eval, used, recheck);
                     case Project project -> inlineJoin ? pruneColumnsInProject(project, used) : p;
                     case EsRelation esr -> pruneColumnsInEsRelation(esr, used);
+                    case Fork fork -> {
+                        forkPresent.set(true);
+                        yield pruneColumnsInFork(fork, used);
+                    }
                     default -> p;
                 };
             } while (recheck.get());
@@ -94,7 +96,7 @@ private static LogicalPlan pruneColumns(LogicalPlan plan, AttributeSet.Builder u
         });
     }
 
-    private static LogicalPlan pruneColumnsInAggregate(Aggregate aggregate, AttributeSet.Builder used, boolean inlineJoin) {
+    static LogicalPlan pruneColumnsInAggregate(Aggregate aggregate, AttributeSet.Builder used, boolean inlineJoin) {
         LogicalPlan p = aggregate;
 
         var remaining = pruneUnusedAndAddReferences(aggregate.aggregates(), used);
@@ -134,7 +136,7 @@ private static LogicalPlan pruneColumnsInAggregate(Aggregate aggregate, Attribut
         return p;
     }
 
-    private static LogicalPlan pruneColumnsInInlineJoinRight(InlineJoin ij, AttributeSet.Builder used, Holder<Boolean> recheck) {
+    static LogicalPlan pruneColumnsInInlineJoinRight(InlineJoin ij, AttributeSet.Builder used, Holder<Boolean> recheck) {
         LogicalPlan p = ij;
 
         used.addAll(ij.references());
@@ -155,7 +157,7 @@ private static LogicalPlan pruneColumnsInInlineJoinRight(InlineJoin ij, Attribut
         return p;
     }
 
-    private static LogicalPlan pruneColumnsInEval(Eval eval, AttributeSet.Builder used, Holder<Boolean> recheck) {
+    static LogicalPlan pruneColumnsInEval(Eval eval, AttributeSet.Builder used, Holder<Boolean> recheck) {
         LogicalPlan p = eval;
 
         var remaining = pruneUnusedAndAddReferences(eval.fields(), used);
@@ -173,7 +175,7 @@ private static LogicalPlan pruneColumnsInEval(Eval eval, AttributeSet.Builder us
     }
 
     // Note: only run when the Project is a descendent of an InlineJoin.
-    private static LogicalPlan pruneColumnsInProject(Project project, AttributeSet.Builder used) {
+    static LogicalPlan pruneColumnsInProject(Project project, AttributeSet.Builder used) {
         LogicalPlan p = project;
 
         var remaining = pruneUnusedAndAddReferences(project.projections(), used);
@@ -184,7 +186,7 @@ private static LogicalPlan pruneColumnsInProject(Project project, AttributeSet.B
         return p;
     }
 
-    private static LogicalPlan pruneColumnsInEsRelation(EsRelation esr, AttributeSet.Builder used) {
+    static LogicalPlan pruneColumnsInEsRelation(EsRelation esr, AttributeSet.Builder used) {
         LogicalPlan p = esr;
 
         if (esr.indexMode() == IndexMode.LOOKUP) {
@@ -200,6 +202,36 @@ private static LogicalPlan pruneColumnsInEsRelation(EsRelation esr, AttributeSet
         return p;
     }
 
+    private static LogicalPlan pruneColumnsInFork(Fork fork, AttributeSet.Builder used) {
+        // prune the output attributes of fork based on usage from the rest of the plan
+        // this does not consider the inner usage within each branch of the fork
+        // as those will be handled when traversing down each branch in PruneColumnsInForkBranches
+        LogicalPlan p = fork;
+
+        // should exit early for UnionAll
+        if (fork instanceof UnionAll) {
+            return p;
+        }
+        boolean changed = false;
+        AttributeSet.Builder builder = AttributeSet.builder();
+        // if any of the fork outputs are used, keep them
+        // otherwise, prune them based on the rest of the plan's usage
+        Set<String> names = new HashSet<>(used.build().names());
+        for (var attr : fork.output()) {
+            // we should also ensure to keep any synthetic attributes around as those could still be used for internal processing
+            if (attr.synthetic() || names.contains(attr.name())) {
+                builder.add(attr);
+            } else {
+                changed = true;
+            }
+        }
+        if (changed) {
+            List<Attribute> attrs = builder.build().stream().toList();
+            p = new Fork(fork.source(), fork.children(), attrs);
+        }
+        return p;
+    }
+
     private static LogicalPlan emptyLocalRelation(UnaryPlan plan) {
         // create an empty local relation with no attributes
         return skipPlan(plan);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumnsInForkBranches.java

@@ -0,0 +1,127 @@
+/*
+ * 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.optimizer.rules.logical;
+
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
+import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
+import org.elasticsearch.xpack.esql.core.util.Holder;
+import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
+import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
+import org.elasticsearch.xpack.esql.plan.logical.Eval;
+import org.elasticsearch.xpack.esql.plan.logical.Fork;
+import org.elasticsearch.xpack.esql.plan.logical.Limit;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.Project;
+import org.elasticsearch.xpack.esql.plan.logical.Sample;
+import org.elasticsearch.xpack.esql.plan.logical.UnionAll;
+import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
+import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
+import org.elasticsearch.xpack.esql.rule.Rule;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns.pruneColumnsInAggregate;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns.pruneColumnsInEsRelation;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns.pruneColumnsInEval;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns.pruneColumnsInInlineJoinRight;
+
+/**
+ * This is used to prune unused columns and expressions in each branch of a Fork.
+ * The output for each fork branch has already been pruned in {@code PruneColumns#pruneColumnsInFork}, so here we only need to
+ * remove unused columns and expressions in the sub-plans of each branch, similarly to independently running {@code PruneColumns}.
+ */
+public final class PruneColumnsInForkBranches extends Rule<LogicalPlan, LogicalPlan> {
+
+    @Override
+    public LogicalPlan apply(LogicalPlan plan) {
+
+        // collect used attributes from the plan above fork
+        var used = AttributeSet.forkBuilder();
+        var forkFound = new Holder<>(false);
+
+        // traverse down to the fork node
+        return plan.transformDown(p -> {
+            // if fork is not found yet, keep collecting used attributes from everything above.
+            // Once fork is found, just return the rest of the plan as is, as any pruning/transformation will have
+            // taken place in pruneSubPlan for each of the fork branches.
+            if (false == (p instanceof Fork) || forkFound.get()) {
+                if (false == forkFound.get()) {
+                    used.addAll(p.references());
+                }
+                return p;
+            }
+
+            // only do this for fork
+            if (p instanceof UnionAll) {
+                return p;
+            }
+
+            used.addAll(p.output());
+            forkFound.set(true);
+            var forkOutputNames = p.output().stream().map(NamedExpression::name).collect(Collectors.toSet());
+            boolean changed = false;
+            List<LogicalPlan> newChildren = new ArrayList<>();
+            for (var child : p.children()) {
+                var clonedUsed = AttributeSet.forkBuilder().addAll(used);
+                var newChild = pruneSubPlan(child, clonedUsed, forkOutputNames);
+                newChildren.add(newChild);
+                if (false == newChild.equals(child)) {
+                    changed = true;
+                }
+            }
+            if (changed) {
+                return new Fork(p.source(), newChildren, p.output());
+            } else {
+                return p;
+            }
+        });
+    }
+
+    private static LogicalPlan pruneSubPlan(LogicalPlan plan, AttributeSet.Builder usedAttrs, Set<String> forkOutput) {
+        if (plan instanceof LocalRelation localRelation) {
+            var outputAttrs = localRelation.output().stream().filter(usedAttrs::contains).collect(Collectors.toList());
+            return new LocalRelation(localRelation.source(), outputAttrs, localRelation.supplier());
+        }
+
+        var projectHolder = new Holder<>(false);
+        return plan.transformDown(p -> {
+            if (p instanceof Limit || p instanceof Sample) {
+                return p;
+            }
+
+            var recheck = new Holder<Boolean>();
+            do {
+                // we operate using the names of the fields, rather than comparing the attributes directly,
+                // as attributes may have been recreated during the transformations of fork branches.
+                recheck.set(false);
+                p = switch (p) {
+                    case Aggregate agg -> pruneColumnsInAggregate(agg, usedAttrs, false);
+                    case InlineJoin inj -> pruneColumnsInInlineJoinRight(inj, usedAttrs, recheck);
+                    case Eval eval -> pruneColumnsInEval(eval, usedAttrs, recheck);
+                    case Project project -> {
+                        // process only the direct Project after Fork, but skip any subsequent instances
+                        if (projectHolder.get()) {
+                            yield p;
+                        } else {
+                            projectHolder.set(true);
+                            var prunedAttrs = project.projections().stream().filter(x -> forkOutput.contains(x.name())).toList();
+                            yield new Project(project.source(), project.child(), prunedAttrs);
+                        }
+                    }
+                    case EsRelation esr -> pruneColumnsInEsRelation(esr, usedAttrs);
+                    default -> p;
+                };
+            } while (recheck.get());
+            usedAttrs.addAll(p.references());
+            return p;
+        });
+    }
+}

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

@@ -162,7 +162,7 @@ public static Set<String> outputUnsupportedAttributeNames(List<LogicalPlan> subp
 
     @Override
     public int hashCode() {
-        return Objects.hash(Fork.class, children());
+        return Objects.hash(Fork.class, output, children());
     }
 
     @Override
@@ -175,7 +175,7 @@ public boolean equals(Object o) {
         }
         Fork other = (Fork) o;
 
-        return Objects.equals(children(), other.children());
+        return Objects.equals(output, other.output) && Objects.equals(children(), other.children());
     }
 
     @Override