ES|QL: fix stats by constant expresson with alias (#117551) (#117613)

Author: luigidellaquila

docs/changelog/117551.yaml

@@ -0,0 +1,5 @@
+pr: 117551
+summary: Fix stats by constant expresson with alias
+area: ES|QL
+type: bug
+issues: []

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

@@ -2758,6 +2758,18 @@ m:integer | y+1:integer
 11        | 12
 ;
 
+statsByConstantExpressionWithAliasAndSort
+required_capability: fix_stats_by_foldable_expression_2
+FROM employees
+| EVAL y = "a"
+| STATS count = COUNT() BY x = y
+| SORT x
+;
+
+count:long | x:keyword
+100        | a
+;
+
 filterIsAlwaysTrue
 required_capability: per_agg_filtering
 FROM employees

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

@@ -504,7 +504,12 @@ public enum Cap {
         /**
          * LOOKUP JOIN
          */
-        JOIN_LOOKUP(Build.current().isSnapshot());
+        JOIN_LOOKUP(Build.current().isSnapshot()),
+
+        /**
+         * Fix for https://github.com/elastic/elasticsearch/issues/114714, again
+         */
+        FIX_STATS_BY_FOLDABLE_EXPRESSION_2,;
 
         private final boolean enabled;
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -463,7 +463,7 @@ static Set<String> fieldNames(LogicalPlan parsed, Set<String> enrichPolicyMatchF
             // remove any already discovered UnresolvedAttributes that are in fact aliases defined later down in the tree
             // for example "from test | eval x = salary | stats max = max(x) by gender"
             // remove the UnresolvedAttribute "x", since that is an Alias defined in "eval"
-            AttributeSet planRefs = Expressions.references(p.expressions());
+            AttributeSet planRefs = p.references();
             p.forEachExpressionDown(Alias.class, alias -> {
                 // do not remove the UnresolvedAttribute that has the same name as its alias, ie "rename id = id"
                 // or the UnresolvedAttributes that are used in Functions that have aliases "STATS id = MAX(id)"

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/session/IndexResolverFieldNamesTests.java

@@ -353,6 +353,114 @@ public void testDocsStats() {
             | SORT languages""", Set.of("emp_no", "emp_no.*", "languages", "languages.*"));
     }
 
+    public void testEvalStats() {
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY y""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY y
+            | SORT y""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY x = y
+            | SORT x""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | STATS count = COUNT(*) BY first_name
+            | SORT first_name""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY x = y
+            | SORT x, first_name""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL first_name = "a"
+            | STATS count = COUNT(*) BY first_name
+            | SORT first_name""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY first_name = to_upper(y)
+            | SORT first_name""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = to_upper(first_name), z = "z"
+            | STATS count = COUNT(*) BY first_name = to_lower(y), z
+            | SORT first_name""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY x = y, z = first_name
+            | SORT x, z""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY x = y, first_name
+            | SORT x, first_name""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(first_name) BY x = y
+            | SORT x
+            | DROP first_name""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY x = y
+            | MV_EXPAND x""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY first_name, y
+            | MV_EXPAND first_name""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | MV_EXPAND first_name
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY first_name, y
+            | SORT y""", Set.of("first_name", "first_name.*"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | MV_EXPAND y
+            | STATS count = COUNT(*) BY x = y
+            | SORT x""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY x = y
+            | STATS count = COUNT(count) by x
+            | SORT x""", Set.of("_index"));
+
+        assertFieldNames("""
+            FROM employees
+            | EVAL y = "a"
+            | STATS count = COUNT(*) BY first_name, y
+            | STATS count = COUNT(count) by x = y
+            | SORT x""", Set.of("first_name", "first_name.*"));
+    }
+
     public void testSortWithLimitOne_DropHeight() {
         assertFieldNames("from employees | sort languages | limit 1 | drop height*", ALL_FIELDS);
     }