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

@@ -4392,3 +4392,20 @@ row a = 1
 c:long
 1
 ;
+
+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

@@ -3463,3 +3463,72 @@ employees:long | job_positions:keyword
 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
+;

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

@@ -1608,6 +1608,13 @@ public enum Cap {
          */
         PUSHING_DOWN_EVAL_WITH_SCORE,
 
+        /**
+         * Fix for ClassCastException in STATS
+         * https://github.com/elastic/elasticsearch/issues/133992
+         * https://github.com/elastic/elasticsearch/issues/136598
+         */
+        FIX_STATS_CLASSCAST_EXCEPTION,
+
         /**
          * Fix attribute equality to respect the name id of the attribute.
          */

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineLimitTopN;
 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.HoistRemoteEnrichLimit;
@@ -178,6 +179,14 @@ protected static Batch<LogicalPlan> operators() {
             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 implements OptimizerRules.CoordinatorOnly {
+
+    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/LocalLogicalPlanOptimizerTests.java

@@ -1656,6 +1656,21 @@ public void testLengthInWhereAndEval() {
     /**
      * Pushed LENGTH to the same field in a <strong>ton</strong> of unique and curious ways. All
      * of these pushdowns should be fused to one.
+     *
+     * <pre>{@code
+     * Project[[l{r}#23]]
+     * \_Eval[[$$SUM$SUM(LENGTH(last>$0{r$}#37 / $$COUNT$$$AVG$SUM(LENGTH(last>$1$1{r$}#41 AS $$AVG$SUM(LENGTH(last>$1#38, $
+     * $SUM$SUM(LENGTH(last>$0{r$}#37 + $$AVG$SUM(LENGTH(last>$1{r$}#38 + $$SUM$SUM(LENGTH(last>$2{r$}#39 AS l#23]]
+     *   \_Limit[1000[INTEGER],false,false]
+     *     \_Aggregate[[],[SUM($$LENGTH(last_nam>$SUM$0{r$}#35,true[BOOLEAN],PT0S[TIME_DURATION],compensated[KEYWORD]) AS $$SUM$SUM(LE
+     *          NGTH(last>$0#37,
+     *          COUNT(a3{r}#11,true[BOOLEAN],PT0S[TIME_DURATION]) AS $$COUNT$$$AVG$SUM(LENGTH(last>$1$1#41,
+     *          SUM($$LENGTH(first_na>$SUM$1{r$}#36,true[BOOLEAN],PT0S[TIME_DURATION],compensated[KEYWORD]) AS $$SUM$SUM(LENGTH(last>$2#39]]
+     *       \_Eval[[$$last_name$LENGTH$920787299{f$}#42 AS a3#11, $$last_name$LENGTH$920787299{f$}#42 AS $$LENGTH(last_nam>$SUM$0
+     * #35, $$first_name$LENGTH$920787299{f$}#43 AS $$LENGTH(first_na>$SUM$1#36]]
+     *         \_Filter[$$last_name$LENGTH$920787299{f$}#42 > 1[INTEGER]]
+     *           \_EsRelation[test][_meta_field{f}#30, emp_no{f}#24, first_name{f}#25, ..]
+     * }</pre>
      */
     public void testLengthPushdownZoo() {
         assumeTrue("requires push", EsqlCapabilities.Cap.VECTOR_SIMILARITY_FUNCTIONS_PUSHDOWN.isEnabled());
@@ -1674,13 +1689,26 @@ public void testLengthPushdownZoo() {
         // Eval - computes final aggregation result (SUM + AVG + SUM)
         var eval1 = as(project.child(), Eval.class);
         assertThat(eval1.fields(), hasSize(2));
+        // The avg is computed as the SUM(LENGTH(last_name)) / COUNT(LENGTH(last_name))
+        var avg = eval1.fields().get(0);
+        var avgDiv = as(avg.child(), Div.class);
+        // SUM(LENGTH(last_name))
+        var evalSumLastName = as(avgDiv.left(), ReferenceAttribute.class);
+        var evalCountLastName = as(avgDiv.right(), ReferenceAttribute.class);
+        var finalAgg = as(eval1.fields().get(1).child(), Add.class);
+        var leftFinalAgg = as(finalAgg.left(), Add.class);
+        assertThat(leftFinalAgg.left(), equalTo(evalSumLastName));
+        assertThat(as(leftFinalAgg.right(), ReferenceAttribute.class).id(), equalTo(avg.id()));
+        // SUM(LENGTH(first_name))
+        var evalSumFirstName = as(finalAgg.right(), ReferenceAttribute.class);
 
         // Limit[1000[INTEGER],false,false]
         var limit = as(eval1.child(), Limit.class);
 
-        // Aggregate with 4 aggregates: SUM for last_name, SUM and COUNT for AVG(a3), SUM for first_name
+        // Aggregate with 3 aggregates: SUM for last_name, COUNT for last_name
+        // (the AVG uses the sum and the count), SUM for first_name
         var agg = as(limit.child(), Aggregate.class);
-        assertThat(agg.aggregates(), hasSize(4));
+        assertThat(agg.aggregates(), hasSize(3));
 
         // Eval - pushdown fields: a3, LENGTH(last_name) for SUM, and LENGTH(first_name) for SUM
         var evalPushdown = as(agg.child(), Eval.class);
@@ -1694,14 +1722,23 @@ public void testLengthPushdownZoo() {
         Attribute firstNamePushDownAttr = assertLengthPushdown(firstNamePushdownAlias.child(), "first_name");
 
         // Verify aggregates reference the pushed down fields
-        var sumForLastName = as(as(agg.aggregates().get(0), Alias.class).child(), Sum.class);
+        var sumForLastNameAlias = as(agg.aggregates().get(0), Alias.class);
+        var sumForLastName = as(sumForLastNameAlias.child(), Sum.class);
         assertThat(as(sumForLastName.field(), ReferenceAttribute.class).id(), equalTo(lastNamePushdownAlias.id()));
-        var sumForAvg = as(as(agg.aggregates().get(1), Alias.class).child(), Sum.class);
-        assertThat(as(sumForAvg.field(), ReferenceAttribute.class).id(), equalTo(a3Alias.id()));
-        var countForAvg = as(as(agg.aggregates().get(2), Alias.class).child(), Count.class);
+        // Checks that the SUM(LENGTH(last_name)) in the final EVAL is the aggregate result here
+        assertThat(evalSumLastName.id(), equalTo(sumForLastNameAlias.id()));
+
+        var countForAvgAlias = as(agg.aggregates().get(1), Alias.class);
+        var countForAvg = as(countForAvgAlias.child(), Count.class);
         assertThat(as(countForAvg.field(), ReferenceAttribute.class).id(), equalTo(a3Alias.id()));
-        var sumForFirstName = as(as(agg.aggregates().get(3), Alias.class).child(), Sum.class);
+        // Checks that the COUNT(LENGTH(last_name)) in the final EVAL is the aggregate result here
+        assertThat(evalCountLastName.id(), equalTo(countForAvgAlias.id()));
+
+        var sumForFirstNameAlias = as(agg.aggregates().get(2), Alias.class);
+        var sumForFirstName = as(sumForFirstNameAlias.child(), Sum.class);
         assertThat(as(sumForFirstName.field(), ReferenceAttribute.class).id(), equalTo(firstNamePushdownAlias.id()));
+        // Checks that the SUM(LENGTH(first_name)) in the final EVAL is the aggregate result here
+        assertThat(evalSumFirstName.id(), equalTo(sumForFirstNameAlias.id()));
 
         // Filter[LENGTH(last_name) > 1]
         var filter = as(evalPushdown.child(), Filter.class);