Fix declared MVs wrongly optimised past STATS

jackpan123 · Dec 5, 2024 · 7842c05 · 7842c05
1 parent 2051401
commit 7842c05
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec
@@ -1095,6 +1095,23 @@ c:long | a:long
 1      | 1
 ;
 
+evalGroupAfterStatsOverRow
+ROW type = ["one", "two"]
+| STATS BY type
+| EVAL is_one = type == "one"
+| EVAL is_two = type == "two"
+;
+
+warning:Line 3:17: evaluation of [type == \"one\"] failed, treating result as null. Only first 20 failures recorded.
+warning:Line 3:17: java.lang.IllegalArgumentException: single-value function encountered multi-value
+warning:Line 4:17: evaluation of [type == \"two\"] failed, treating result as null. Only first 20 failures recorded.
+warning:Line 4:17: java.lang.IllegalArgumentException: single-value function encountered multi-value
+
+type:keyword   |is_one:boolean |is_two:boolean
+one            |true           |false
+two            |false          |true
+;
+
 docsGettingStartedStats
 // tag::gs-stats[]
 FROM sample_data

diff --git a/...ain/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PropagateEvalFoldables.java b/...ain/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PropagateEvalFoldables.java
@@ -12,11 +12,14 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
+import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.rule.Rule;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 /**
  * Replace any reference attribute with its source, if it does not affect the result.
  * This avoids ulterior look-ups between attributes and its source across nodes.
@@ -51,7 +54,18 @@ public LogicalPlan apply(LogicalPlan plan) {
             // TODO: also allow aggregates once aggs on constants are supported.
             // C.f. https://github.com/elastic/elasticsearch/issues/100634
             if (p instanceof Filter || p instanceof Eval) {
-                p = p.transformExpressionsOnly(ReferenceAttribute.class, replaceReference);
+                // Stop propagation if an aggregate is found in the children
+                AtomicBoolean findAggregatePlan = new AtomicBoolean(false);
+                p.transformDown(childPlan -> {
+                    if (childPlan instanceof Aggregate) {
+                        findAggregatePlan.set(true);
+                        return childPlan;
+                    }
+                    return childPlan;
+                });
+                if (findAggregatePlan.get() == false) {
+                    p = p.transformExpressionsOnly(ReferenceAttribute.class, replaceReference);
+                }
             }
             return p;
         });