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

docs/changelog/137511.yaml

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

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,74 @@ VALUES(color):keyword | color:keyword
                yellow | yellow
 // end::mv-group-values-expand-result[]
 ;
+
+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

@@ -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);
     }