iter

Author: pmpailis

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Count.java

@@ -146,7 +146,7 @@ public Nullability nullable() {
     public boolean isWildcard() {
         Expression field = field();
         if (field instanceof Literal lit) {
-            return StringUtils.WILDCARD.equals(lit.value());
+            return StringUtils.WILDCARD.equals(lit.toString());
         }
         return false;
     }

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

@@ -8,6 +8,7 @@
 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;
@@ -24,13 +25,13 @@
 
 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;
-import static org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns.pruneColumnsInProject;
 
 /**
  * This is used to prune unused columns and expressions in each branch of a Fork.
@@ -65,12 +66,12 @@ public LogicalPlan apply(LogicalPlan plan) {
 
             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);
+                var newChild = pruneSubPlan(child, clonedUsed, forkOutputNames);
                 newChildren.add(newChild);
                 if (false == newChild.equals(child)) {
                     changed = true;
@@ -84,7 +85,7 @@ public LogicalPlan apply(LogicalPlan plan) {
         });
     }
 
-    private static LogicalPlan pruneSubPlan(LogicalPlan plan, AttributeSet.Builder usedAttrs) {
+    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());
@@ -111,7 +112,8 @@ private static LogicalPlan pruneSubPlan(LogicalPlan plan, AttributeSet.Builder u
                             yield p;
                         } else {
                             projectHolder.set(true);
-                            yield pruneColumnsInProject(project, usedAttrs);
+                            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);

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java

@@ -70,7 +70,6 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToString;
 import org.elasticsearch.xpack.esql.expression.function.scalar.internal.PackDimension;
 import org.elasticsearch.xpack.esql.expression.function.scalar.internal.UnpackDimension;
-import org.elasticsearch.xpack.esql.expression.function.scalar.math.Exp;
 import org.elasticsearch.xpack.esql.expression.function.scalar.math.Round;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvAvg;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvCount;
@@ -9643,14 +9642,18 @@ public void testPruneColumnsInForkBranchesDropNestedEvalsIfNotNeeded() {
             var filter = as(limitInBranch.child(), Filter.class);
             assertThat(filter.toString(), containsString("IN(10048[INTEGER],10081[INTEGER],emp_no"));
             var relation = as(filter.child(), EsRelation.class);
-            assertThat(relation.output().stream().map(Attribute::name).collect(Collectors.toSet()), hasItems("emp_no", "first_name", "last_name"));
+            assertThat(
+                relation.output().stream().map(Attribute::name).collect(Collectors.toSet()),
+                hasItems("emp_no", "first_name", "last_name")
+            );
         }
     }
 
     /**
      * Limit[10000[INTEGER],false,false]
      * \_Aggregate[[],[COUNT(*[KEYWORD],true[BOOLEAN],PT0S[TIME_DURATION]) AS count(*)#36, COUNT(emp_no{r}#109,true[BOOLEAN],PT0S[
-     * TIME_DURATION]) AS d#39, MAX(_fork{r}#124,true[BOOLEAN],PT0S[TIME_DURATION]) AS m#42, COUNT(s{r}#119,true[BOOLEAN],PT0S[TIME_DURATION]) AS ls#45]]
+     * TIME_DURATION]) AS d#39, MAX(_fork{r}#124,true[BOOLEAN],PT0S[TIME_DURATION]) AS m#42,
+     * COUNT(s{r}#119,true[BOOLEAN],PT0S[TIME_DURATION]) AS ls#45]]
      *   \_Fork[[emp_no{r}#109, s{r}#119, _fork{r}#124]]
      *     |_Project[[emp_no{f}#46, s{r}#6, _fork{r}#7]]
      *     | \_Eval[[fork1[KEYWORD] AS _fork#7]]
@@ -9722,7 +9725,10 @@ public void testPruneColumnsInForkBranchesDoNotPruneIfWildcardAggregation() {
         var evalInSecondBranch = as(secondBranchProject.child(), Eval.class);
         var limitInSecondBranch = as(evalInSecondBranch.child(), Limit.class);
         var localRelation = as(limitInSecondBranch.child(), LocalRelation.class);
-        assertThat(localRelation.output().stream().map(Attribute::name).collect(Collectors.toSet()), hasItems("$$COUNT$COUNT(*)::keywo>$0"));
+        assertThat(
+            localRelation.output().stream().map(Attribute::name).collect(Collectors.toSet()),
+            hasItems("$$COUNT$COUNT(*)::keywo>$0")
+        );
 
         var thirdBranch = fork.children().get(2);
         var thirdBranchProject = as(thirdBranch, Project.class);