Mirror upstream elastic#137511 as single snapshot commit for AI review

BASE=3017e334274a7292997b0fea77f90d2c73b58eba
HEAD=71c1691e34614a43230da1905611f4fa9dfeec01
Branch=main

Author: phananh1010

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/stats.csv-spec

@@ -3463,3 +3463,89 @@ 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
+;
+
+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
+;
+
+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
+;
+
+fixClassCastBug6
+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

@@ -1579,6 +1579,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,8 @@ protected static Batch<LogicalPlan> operators() {
             new SplitInWithFoldableValue(),
             new PropagateEvalFoldables(),
             new ConstantFolding(),
+            // then extract nested aggs top-level
+            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,103 @@
+/*
+ * 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;
+
+import org.elasticsearch.common.util.Maps;
+import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.util.Holder;
+import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
+import org.elasticsearch.xpack.esql.expression.function.grouping.GroupingFunction;
+import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.Project;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * 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 OptimizerRules.OptimizerRule<Aggregate> implements OptimizerRules.CoordinatorOnly {
+    public DeduplicateAggs() {
+        super(OptimizerRules.TransformDirection.UP);
+    }
+
+    @Override
+    protected LogicalPlan rule(Aggregate aggregate) {
+        // an alias map for evaluatable grouping functions
+        AttributeMap.Builder<Expression> aliasesBuilder = AttributeMap.builder();
+        aggregate.forEachExpressionUp(Alias.class, a -> {
+            if (a.child() instanceof GroupingFunction.NonEvaluatableGroupingFunction == false) {
+                aliasesBuilder.put(a.toAttribute(), a.child());
+            }
+        });
+        var aliases = aliasesBuilder.build();
+
+        // break down each aggregate into AggregateFunction and/or grouping key
+        // preserve the projection at the end
+        List<? extends NamedExpression> aggs = aggregate.aggregates();
+
+        // root/naked aggs
+        Map<AggregateFunction, Alias> rootAggs = Maps.newLinkedHashMapWithExpectedSize(aggs.size());
+        List<NamedExpression> newProjections = new ArrayList<>();
+        // track the aggregate aggs (including grouping which is not an AggregateFunction)
+        List<NamedExpression> newAggs = new ArrayList<>();
+
+        Holder<Boolean> changed = new Holder<>(false);
+
+        for (NamedExpression agg : aggs) {
+            if (agg instanceof Alias as) {
+                // use intermediate variable to mark child as final for lambda use
+                Expression child = as.child();
+
+                // common case - handle duplicates
+                if (child instanceof AggregateFunction af) {
+                    // canonical representation, with resolved aliases
+                    AggregateFunction canonical = (AggregateFunction) af.transformUp(e -> aliases.resolve(e, e));
+
+                    Alias found = rootAggs.get(canonical);
+                    // aggregate is new
+                    if (found == null) {
+                        rootAggs.put(canonical, as);
+                        newAggs.add(as);
+                        newProjections.add(as.toAttribute());
+                    }
+                    // agg already exists - preserve the current alias but point it to the existing agg
+                    // thus don't add it to the list of aggs as we don't want duplicated compute
+                    else {
+                        changed.set(true);
+                        newProjections.add(as.replaceChild(found.toAttribute()));
+                    }
+                }
+            }
+            // not an alias (e.g. grouping field)
+            else {
+                newAggs.add(agg);
+                newProjections.add(agg.toAttribute());
+            }
+        }
+
+        LogicalPlan plan = aggregate;
+        if (changed.get()) {
+            Source source = aggregate.source();
+            plan = aggregate.with(aggregate.child(), aggregate.groupings(), newAggs);
+            // preserve initial projection
+            plan = new Project(source, plan, newProjections);
+        }
+
+        return plan;
+    }
+}

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

@@ -109,11 +109,12 @@ protected LogicalPlan rule(Aggregate aggregate) {
                 else {
                     changed.set(true);
                     Expression aggExpression = child.transformUp(AggregateFunction.class, af -> {
-                        AggregateFunction canonical = (AggregateFunction) af.canonical();
+                        // canonical representation, with resolved aliases
+                        AggregateFunction canonical = (AggregateFunction) af.canonical().transformUp(e -> aliases.resolve(e, e));
                         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

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

@@ -533,30 +533,6 @@ public void testAggsWithOverridingInputAndGrouping() throws Exception {
         );
     }
 
-    /**
-     * <pre>{@code
-     * 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]]
-     *     \_EsRelation[test][_meta_field{f}#23, emp_no{f}#17, first_name{f}#18, ..]
-     * }</pre>
-     */
-    public void testCombineProjectionWithDuplicateAggregation() {
-        var plan = plan("""
-            from test
-            | stats s = sum(salary), d = sum(salary), c = sum(salary) by last_name, first_name
-            | keep d, s, last_name, first_name
-            """);
-
-        var project = as(plan, Project.class);
-        assertThat(Expressions.names(project.projections()), contains("d", "s", "last_name", "first_name"));
-        var limit = as(project.child(), Limit.class);
-        var agg = as(limit.child(), Aggregate.class);
-        assertThat(Expressions.names(agg.aggregates()), contains("s", "last_name", "first_name"));
-        assertThat(Alias.unwrap(agg.aggregates().get(0)), instanceOf(Sum.class));
-        assertThat(Expressions.names(agg.groupings()), contains("last_name", "first_name"));
-    }
-
     /**
      * <pre>{@code
      * Limit[1000[INTEGER]]