Fixes esql class cast bug in STATS at planning level (#137511)

Author: ncordon

docs/changelog/137511.yaml

@@ -0,0 +1,5 @@
+pr: 137511
+summary: Fixes esql class cast bug in STATS at planning level
+area: ES|QL
+type: bug
+issues: []

x-pack/plugin/esql/qa/testFixtures/src/main/resources/inlinestats.csv-spec

@@ -1001,3 +1001,20 @@ FROM hosts METADATA _index
   description:text |  host:keyword |   ip0:ip      | _index:keyword |   x:ip  | ip1:long | host_group:text | card:keyword
 alpha db server    |alpha          |127.0.0.1      |hosts           |127.0.0.1|1         |DB servers       |eth0
 ;
+
+fixClassCastBugWithSeveralCounts
+required_capability: inline_stats
+required_capability: fix_stats_classcast_exception
+
+FROM sample_data, sample_data_str
+| EVAL one_ip = client_ip::ip
+| INLINE STATS count1=count(client_ip::ip), count2=count(one_ip)
+| KEEP count1, count2
+| LIMIT 3
+;
+
+count1:long    |count2:long
+14             |14 
+14             |14 
+14             |14
+;

x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec

@@ -3344,3 +3344,103 @@ VALUES(color):keyword | color:keyword
                yellow | yellow
 // end::mv-group-values-expand-result[]
 ;
+<<<<<<< HEAD
+=======
+
+filterOrdinalValues
+required_capability: fix_filter_ordinals
+from employees 
+| MV_EXPAND job_positions
+| WHERE job_positions != "Purchase Manager"
+| STATS employees=count(*) BY job_positions
+| SORT job_positions
+| LIMIT 1000
+;
+
+employees:long | job_positions:keyword
+18             | Accountant
+13             | Architect
+11             | Business Analyst
+13             | Data Scientist
+10             | Head Human Resources
+15             | Internship
+14             | Junior Developer
+20             | Principal Support Engineer
+13             | Python Developer
+15             | Reporting Analyst
+20             | Senior Python Developer
+15             | Senior Team Lead
+11             | Support Engineer
+15             | Tech Lead    
+;
+
+fixClassCastBugWithCountDistinct
+required_capability: fix_stats_classcast_exception
+
+from airports 
+| rename scalerank AS x 
+| stats  a = count(x), b = count(x) + count(x), c = count_distinct(x)
+;
+
+a:long | b:long | c:long
+891    | 1782   | 8
+;
+
+fixClassCastBugWithValuesFn
+required_capability: fix_stats_classcast_exception
+
+ROW x = [1,2,3] 
+| STATS a = MV_COUNT(VALUES(x)), b = VALUES(x), c = SUM(x)
+;
+
+a:integer | b:integer | c:long
+3         | [1, 2, 3] | 6
+;
+
+fixClassCastBugWithSeveralCountDistincts
+required_capability: fix_stats_classcast_exception
+
+ROW x = 1 
+| STATS a = 2*COUNT_DISTINCT(x), b = COUNT_DISTINCT(x), c = MAX(x)
+;
+
+a:long | b:long | c:integer
+2      | 1      | 1
+;
+
+fixClassCastBugWithMedianPlusCountDistinct
+required_capability: fix_stats_classcast_exception
+
+FROM sample_data_ts_long 
+| EVAL sym1 = 0, sym5 = 1
+| STATS sym2 = median(sym5) + 0, sym3 = median(sym5), sym4 = count_distinct(sym1)
+;
+
+sym2:double |sym3:double | sym4:long
+1.0         | 1.0        | 1
+;
+
+fixClassCastBugWithFoldableLiterals
+required_capability: fix_stats_classcast_exception
+
+from airports 
+| rename scalerank AS x 
+| stats  a = count(x), b = count(x) + count(x), c = count_distinct(x, 10), d = count_distinct(x, 10 + 1 - 1)
+;
+
+a:long | b:long | c:long | d:long
+891    | 1782   | 8      | 8
+;
+
+fixClassCastBugWithSurrogateExpressions
+required_capability: fix_stats_classcast_exception
+
+from airports 
+| rename scalerank AS x 
+| stats a = median(x), b = percentile(x, 50), c = count_distinct(x)
+;
+
+a:double | b:double | c:long
+6.0      | 6.0      | 8
+;
+>>>>>>> 0a81875703ed (Fixes esql class cast bug in STATS at planning level (#137511))

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1282,6 +1282,13 @@ public enum Cap {
          */
         STATS_WITH_FILTERED_SURROGATE_FIXED,
 
+        /**
+         * Fix for ClassCastException in STATS
+         * https://github.com/elastic/elasticsearch/issues/133992
+         * https://github.com/elastic/elasticsearch/issues/136598
+         */
+        FIX_STATS_CLASSCAST_EXCEPTION,
+
         /**
          * {@link org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAliasingEvalWithProject} did not fully account for shadowing.
          * https://github.com/elastic/elasticsearch/issues/137019.

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineEvals;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineProjections;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ConstantFolding;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.DeduplicateAggs;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ExtractAggregateCommonFilter;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.FoldNull;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.LiteralsOnTheRight;
@@ -171,6 +172,14 @@ protected static Batch<LogicalPlan> operators(boolean local) {
             new SplitInWithFoldableValue(),
             new PropagateEvalFoldables(),
             new ConstantFolding(),
+            /* Then deduplicate aggregations
+               We need this after the constant folding
+               because we could have expressions like
+                    count_distinct(_, 9 + 1)
+                    count_distinct(_, 10)
+               which are semantically identical
+             */
+            new DeduplicateAggs(),
             new PartiallyFoldCase(),
             // boolean
             new BooleanSimplification(),

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

@@ -0,0 +1,21 @@
+/*
+ * 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;
+
+/**
+ * This rule handles duplicate aggregate functions to avoid duplicate compute
+ * stats a = min(x), b = min(x), c = count(*), d = count() by g
+ * becomes
+ * stats a = min(x), c = count(*) by g | eval b = a, d = c | keep a, b, c, d, g
+ */
+public final class DeduplicateAggs extends ReplaceAggregateAggExpressionWithEval {
+
+    public DeduplicateAggs() {
+        super(false);
+    }
+}

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

@@ -43,9 +43,17 @@
  * becomes
  * stats a = min(x), c = count(*) by g | eval b = a, d = c | keep a, b, c, d, g
  */
-public final class ReplaceAggregateAggExpressionWithEval extends OptimizerRules.OptimizerRule<Aggregate> {
+public class ReplaceAggregateAggExpressionWithEval extends OptimizerRules.OptimizerRule<Aggregate> {
+    private final boolean replaceNestedExpressions;
+
+    public ReplaceAggregateAggExpressionWithEval(boolean replaceNestedExpressions) {
+        super(OptimizerRules.TransformDirection.UP);
+        this.replaceNestedExpressions = replaceNestedExpressions;
+    }
+
     public ReplaceAggregateAggExpressionWithEval() {
         super(OptimizerRules.TransformDirection.UP);
+        this.replaceNestedExpressions = true;
     }
 
     @Override
@@ -88,7 +96,7 @@ protected LogicalPlan rule(Aggregate aggregate) {
                 // common case - handle duplicates
                 if (child instanceof AggregateFunction af) {
                     // canonical representation, with resolved aliases
-                    AggregateFunction canonical = (AggregateFunction) af.canonical().transformUp(e -> aliases.resolve(e, e));
+                    AggregateFunction canonical = getCannonical(af, aliases);
 
                     Alias found = rootAggs.get(canonical);
                     // aggregate is new
@@ -106,14 +114,15 @@ protected LogicalPlan rule(Aggregate aggregate) {
                 }
                 // nested expression over aggregate function or groups
                 // replace them with reference and move the expression into a follow-up eval
-                else {
+                else if (replaceNestedExpressions) {
                     changed.set(true);
                     Expression aggExpression = child.transformUp(AggregateFunction.class, af -> {
-                        AggregateFunction canonical = (AggregateFunction) af.canonical();
+                        // canonical representation, with resolved aliases
+                        AggregateFunction canonical = getCannonical(af, aliases);
                         Alias alias = rootAggs.get(canonical);
                         if (alias == null) {
-                            // create synthetic alias ove the found agg function
-                            alias = new Alias(af.source(), syntheticName(canonical, child, counter[0]++), canonical, null, true);
+                            // create synthetic alias over the found agg function
+                            alias = new Alias(af.source(), syntheticName(canonical, child, counter[0]++), af.canonical(), null, true);
                             // and remember it to remove duplicates
                             rootAggs.put(canonical, alias);
                             // add it to the list of aggregates and continue
@@ -132,6 +141,9 @@ protected LogicalPlan rule(Aggregate aggregate) {
                     Alias alias = as.replaceChild(aggExpression);
                     newEvals.add(alias);
                     newProjections.add(alias.toAttribute());
+                } else {
+                    newAggs.add(agg);
+                    newProjections.add(agg.toAttribute());
                 }
             }
             // not an alias (e.g. grouping field)
@@ -155,6 +167,10 @@ protected LogicalPlan rule(Aggregate aggregate) {
         return plan;
     }
 
+    private static AggregateFunction getCannonical(AggregateFunction af, AttributeMap<Expression> aliases) {
+        return (AggregateFunction) af.canonical().transformUp(e -> aliases.resolve(e, e));
+    }
+
     private static String syntheticName(Expression expression, Expression af, int counter) {
         return TemporaryNameUtils.temporaryName(expression, af, counter);
     }

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

@@ -178,6 +178,13 @@
 import static org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison.BinaryComparisonOperation.GTE;
 import static org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison.BinaryComparisonOperation.LT;
 import static org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison.BinaryComparisonOperation.LTE;
+<<<<<<< HEAD
+=======
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.DeduplicateAggsTests.aggFieldName;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.DeduplicateAggsTests.aliased;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.DOWN;
+import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.UP;
+>>>>>>> 0a81875703ed (Fixes esql class cast bug in STATS at planning level (#137511))
 import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.contains;
@@ -473,6 +480,7 @@ public void testAggsWithOverridingInputAndGrouping() throws Exception {
     }
 
     /**
+<<<<<<< HEAD
      * Project[[s{r}#4 AS d, s{r}#4, last_name{f}#21, first_name{f}#18]]
      * \_Limit[1000[INTEGER]]
      *   \_Aggregate[[last_name{f}#21, first_name{f}#18],[SUM(salary{f}#22) AS s, last_name{f}#21, first_name{f}#18]]
@@ -495,6 +503,9 @@ public void testCombineProjectionWithDuplicateAggregation() {
     }
 
     /**
+=======
+     * <pre>{@code
+>>>>>>> 0a81875703ed (Fixes esql class cast bug in STATS at planning level (#137511))
      * Limit[1000[INTEGER]]
      * \_Aggregate[STANDARD,[],[SUM(salary{f}#12,true[BOOLEAN]) AS sum(salary), SUM(salary{f}#12,last_name{f}#11 == [44 6f 65][KEYW
      * ORD]) AS sum(salary) WheRe last_name ==   "Doe"]]
@@ -3809,6 +3820,7 @@ public void testPruneRenameOnAggBy() {
     }
 
     /**
+<<<<<<< HEAD
      * Expects
      * Project[[c1{r}#2, c2{r}#4, cs{r}#6, cm{r}#8, cexp{r}#10]]
      * \_Eval[[c1{r}#2 AS c2, c1{r}#2 AS cs, c1{r}#2 AS cm, c1{r}#2 AS cexp]]
@@ -3971,6 +3983,8 @@ public void testEliminateDuplicateRenamedGroupings() {
     }
 
     /**
+=======
+>>>>>>> 0a81875703ed (Fixes esql class cast bug in STATS at planning level (#137511))
      * Expected
      * Limit[2[INTEGER]]
      * \_Filter[a{r}#6 > 2[INTEGER]]
@@ -4061,19 +4075,6 @@ public void testLimitZeroUsesLocalRelation() {
         assertThat(plan, instanceOf(LocalRelation.class));
     }
 
-    private <T> T aliased(Expression exp, Class<T> clazz) {
-        var alias = as(exp, Alias.class);
-        return as(alias.child(), clazz);
-    }
-
-    private <T extends AggregateFunction> void aggFieldName(Expression exp, Class<T> aggType, String fieldName) {
-        var alias = as(exp, Alias.class);
-        var af = as(alias.child(), aggType);
-        var field = af.field();
-        var name = field.foldable() ? BytesRefs.toString(field.fold(FoldContext.small())) : Expressions.name(field);
-        assertThat(name, is(fieldName));
-    }
-
     /**
      * Expects
      * Limit[1000[INTEGER]]
@@ -4831,6 +4832,7 @@ public void testStatsExpOverAggsWithScalars() {
 
     /**
      * Expects
+<<<<<<< HEAD
      * Project[[a{r}#5, b{r}#9, $$max(salary)_+_3>$COUNT$2{r}#46 AS d, $$count(salary)_->$MIN$3{r}#47 AS e, $$avg(salary)_+_m
      * >$MAX$1{r}#45 AS g]]
      * \_Eval[[$$$$avg(salary)_+_m>$AVG$0$SUM$0{r}#48 / $$max(salary)_+_3>$COUNT$2{r}#46 AS $$avg(salary)_+_m>$AVG$0, $$avg(
@@ -4902,6 +4904,9 @@ public void testStatsExpOverAggsWithScalarAndDuplicateAggs() {
 
     /**
      * Expects
+=======
+     * <pre>{@code
+>>>>>>> 0a81875703ed (Fixes esql class cast bug in STATS at planning level (#137511))
      * Project[[a{r}#5, a{r}#5 AS b, w{r}#12]]
      * \_Limit[1000[INTEGER]]
      *   \_Aggregate[[w{r}#12],[SUM($$salary_/_2_+_la>$SUM$0{r}#26) AS a, w{r}#12]]