ESQL: Fix columns ordering when pruning an INLINE STATS (#136827)

This fixes the incorrect ordering of the columns that can occur when
INLINE STATS is pruned entirely, due to its calculated value being
discarded (overshadowed).

Fixes #136797

Author: bpintea

docs/changelog/136827.yaml

@@ -0,0 +1,6 @@
+pr: 136827
+summary: Fix columns ordering when pruning an INLINE STATS
+area: ES|QL
+type: bug
+issues:
+ - 136797

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java

@@ -71,7 +71,6 @@ public abstract class GenerativeRestTest extends ESRestTestCase implements Query
         "cannot be cast to class", // https://github.com/elastic/elasticsearch/issues/133992
         "can't find input for", // https://github.com/elastic/elasticsearch/issues/136596
         "unexpected byte", // https://github.com/elastic/elasticsearch/issues/136598
-        "Output has changed from", // https://github.com/elastic/elasticsearch/issues/136797
         "out of bounds for length", // https://github.com/elastic/elasticsearch/issues/136851
 
         // Awaiting fixes for correctness

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

@@ -1662,6 +1662,111 @@ ROW salary = 12300, emp_no = 5, gender = "F"
 emp_no:integer
 5
 ;
+
+inlineStatsWithShadowedOutput1
+required_capability: inline_stats_prune_column_fix
+
+ROW a = null, b = 0, c = "x"
+| INLINE STATS a = MAX(c) BY b
+| EVAL a = c
+;
+
+c:keyword      |b:integer      |a:keyword
+x              |0              |x
+;
+
+inlineStatsWithShadowedOutput2
+required_capability: inline_stats_prune_column_fix
+
+ROW a = null, b = 0, c = "x", d = 1
+| INLINE STATS a = MAX(c), d = MIN(c) by b
+| EVAL a = c, d = 2
+;
+
+c:keyword      |b:integer      |a:keyword |d:integer
+x              |0              |x         |2
+;
+
+inlineStatsWithShadowedOutput3
+required_capability: inline_stats_prune_column_fix
+
+FROM employees
+| KEEP emp_no, salary, birth_date
+| EVAL orig_emp_no = emp_no
+| INLINE STATS emp_no = MAX(birth_date) BY salary
+| EVAL emp_no = birth_date
+| SORT orig_emp_no
+| LIMIT 3
+;
+
+birth_date:date         |orig_emp_no:integer|salary:integer |emp_no:date
+1953-09-02T00:00:00.000Z|10001              |57305          |1953-09-02T00:00:00.000Z
+1964-06-02T00:00:00.000Z|10002              |56371          |1964-06-02T00:00:00.000Z
+1959-12-03T00:00:00.000Z|10003              |61805          |1959-12-03T00:00:00.000Z
+;
+
+inlineStatsWithShadowedOutput4
+required_capability: inline_stats_prune_column_fix
+
+ROW a = null, b = 0, c = "x", d = 1
+| INLINE STATS a = MAX(c), d = MIN(c), e = AVG(b) by b
+| EVAL a = c, d = 2, e = a
+;
+
+c:keyword      |b:integer      |a:keyword      |d:integer      |e:keyword
+x              |0              |x              |2              |x
+;
+
+inlineStatsWithShadowedOutput5
+required_capability: inline_stats_prune_column_fix
+
+ROW a = null, b = 0, c = "x", d = 1
+| INLINE STATS a = MAX(c), d = MIN(c), e = AVG(b), a = MIN(b) by b
+| EVAL a = c, d = 2, e = a
+;
+warning:Line 2:16: Field 'a' shadowed by field at line 2:52
+
+c:keyword      |b:integer      |a:keyword      |d:integer      |e:keyword
+x              |0              |x              |2              |x
+;
+
+inlineStatsWithShadowedOutput6
+required_capability: inline_stats_prune_column_fix
+
+ROW a = null, b = 0, c = "x", d = 1
+| INLINE STATS a = MAX(c), d = MIN(c), e = AVG(b), a = MIN(b), f = MAX(b) by b
+| EVAL a = c, d = 2, e = a, f = d
+;
+warning:Line 2:16: Field 'a' shadowed by field at line 2:52
+
+c:keyword      |b:integer      |a:keyword      |d:integer      |e:keyword      |f:integer
+x              |0              |x              |2              |x              |2
+;
+
+tripleInlineStatsMultipleAssignmentsGetsPrunedPartially
+required_capability: inline_stats_prune_column_fix
+
+FROM employees
+// | SORT languages DESC
+| EVAL salaryK = salary / 10000
+| INLINE STATS count = COUNT(*) BY salaryK
+| EVAL hire_date_string = hire_date::keyword
+| INLINE STATS sum = SUM(languages) BY hire_date_string
+| DISSECT hire_date_string "%{date}"
+| INLINE STATS min = MIN(MV_COUNT(languages)) BY salaryK
+| SORT emp_no
+| KEEP emp_no, salaryK, count, min
+| LIMIT 5
+;
+
+emp_no:integer |salaryK:integer|count:long     |min:integer
+10001          |5              |17             |1
+10002          |5              |17             |1
+10003          |6              |17             |1
+10004          |3              |27             |1
+10005          |6              |17             |1
+;
+
 ///////////////////////////
 // INLINE STATS with filters
 ///////////////////////////

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

@@ -1531,6 +1531,11 @@ public enum Cap {
          */
         GROK_MULTI_PATTERN,
 
+        /**
+         * Fix pruning of columns when shadowed in INLINE STATS
+         */
+        INLINE_STATS_PRUNE_COLUMN_FIX(INLINESTATS.enabled),
+
         /**
          * Fix double release in inline stats when LocalRelation is reused
          */

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

@@ -123,7 +123,7 @@ private static LogicalPlan pruneColumnsInAggregate(Aggregate aggregate, Attribut
         } else {
             // not expecting high groups cardinality, nested loops in lists should be fine, no need for a HashSet
             if (inlineJoin && aggregate.groupings().containsAll(remaining)) {
-                // An INLINEJOIN right-hand side aggregation output had everything pruned, except for (some of the) groupings, which are
+                // An InlineJoin right-hand side aggregation output had everything pruned, except for (some of the) groupings, which are
                 // already part of the IJ output (from the left-hand side): the agg can just be dropped entirely.
                 p = emptyLocalRelation(aggregate);
             } else { // not an INLINEJOIN or there are actually aggregates to compute
@@ -140,7 +140,12 @@ private static LogicalPlan pruneColumnsInInlineJoinRight(InlineJoin ij, Attribut
         used.addAll(ij.references());
         var right = pruneColumns(ij.right(), used, true);
         if (right.output().isEmpty() || isLocalEmptyRelation(right)) {
-            p = ij.left();
+            // InlineJoin updates the order of the output, so even if the computation is dropped, the groups need to be pulled to the end.
+            // So we keep just the left side of the join (i.e. drop the computations), but place a Project on top to keep the right order.
+            List<Attribute> newOutput = new ArrayList<>(ij.output());
+            AttributeSet leftOutputSet = ij.left().outputSet();
+            newOutput.removeIf(attr -> leftOutputSet.contains(attr) == false);
+            p = new Project(ij.source(), ij.left(), newOutput);
             recheck.set(true);
         } else if (right != ij.right()) {
             // if the right side has been updated, replace it

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

@@ -5955,7 +5955,7 @@ public void testDoubleInlineStatsWithEvalGetsPrunedEntirely() {
         }
         var plan = optimizedPlan(query);
 
-        var project = as(plan, EsqlProject.class);
+        var project = as(plan, Project.class);
         assertThat(Expressions.names(project.projections()), is(List.of("emp_no", "count")));
         var topN = as(project.child(), TopN.class);
         assertThat(topN.order().size(), is(1));
@@ -6070,22 +6070,26 @@ public void testTripleInlineStatsGetsPrunedPartially() {
     }
 
     /*
-     * EsqlProject[[emp_no{f}#864, salaryK{r}#836, count{r}#838, min{r}#852]]
-     * \_TopN[[Order[emp_no{f}#864,ASC,LAST]],5[INTEGER]]
-     *   \_InlineJoin[LEFT,[salaryK{r}#836],[salaryK{r}#836],[salaryK{r}#836]]
-     *     |_Dissect[hire_date_string{r}#842,Parser[pattern=%{date}, appendSeparator=,
-     *          parser=org.elasticsearch.dissect.DissectParser@27d57d5e],[date{r}#847]] <-- TODO: Dissect & Eval could/should be dropped
-     *     | \_Eval[[TOSTRING(hire_date{f}#865) AS hire_date_string#842]]
-     *     |   \_InlineJoin[LEFT,[salaryK{r}#836],[salaryK{r}#836],[salaryK{r}#836]]
-     *     |     |_Eval[[salary{f}#866 / 10000[INTEGER] AS salaryK#836]]
-     *     |     | \_EsRelation[employees][emp_no{f}#864, hire_date{f}#865, languages{f}#860, ..]
-     *     |     \_Aggregate[[salaryK{r}#836],[COUNT(*[KEYWORD],true[BOOLEAN]) AS count#838, salaryK{r}#836]]
-     *     |       \_StubRelation[[emp_no{f}#864, hire_date{f}#865, languages{f}#860, languages.byte{f}#861, languages.long{f}#863,
-     *                  languages.short{f}#862, salary{f}#866, salaryK{r}#836]]
-     *     \_Aggregate[[salaryK{r}#836],[MIN($$MV_COUNT(langua>$MIN$0{r$}#867,true[BOOLEAN]) AS min#852, salaryK{r}#836]]
-     *       \_Eval[[MVCOUNT(languages{f}#860) AS $$MV_COUNT(langua>$MIN$0#867]]
-     *         \_StubRelation[[emp_no{f}#864, hire_date{f}#865, languages{f}#860, languages.byte{f}#861, languages.long{f}#863,
-     *              languages.short{f}#862, salary{f}#866, count{r}#838, salaryK{r}#836, sum{r}#845, hire_date_string{r}#842, date{r}#847]]
+     * EsqlProject[[emp_no{f}#26, salaryK{r}#4, count{r}#6, min{r}#19]]
+     * \_TopN[[Order[emp_no{f}#26,ASC,LAST]],5[INTEGER]]
+     *   \_InlineJoin[LEFT,[salaryK{r}#4],[salaryK{r}#4]]
+     *     |_EsqlProject[[_meta_field{f}#32, emp_no{f}#26, first_name{f}#27, gender{f}#28, hire_date{f}#33, job{f}#34, job.raw{f}#35,
+     *              languages{f}#29, last_name{f}#30, long_noidx{f}#36, salary{f}#31, count{r}#6, salaryK{r}#4, hire_date_string{r}#10,
+     *              date{r}#15]]
+     *     | \_Dissect[hire_date_string{r}#10,Parser[pattern=%{date}, appendSeparator=,
+     *            parser=org.elasticsearch.dissect.DissectParser@77d1afc3],[date{r}#15]] <-- TODO: Dissect & Eval could/should be dropped
+     *     |   \_Eval[[TOSTRING(hire_date{f}#33) AS hire_date_string#10]]
+     *     |     \_InlineJoin[LEFT,[salaryK{r}#4],[salaryK{r}#4]]
+     *     |       |_Eval[[salary{f}#31 / 10000[INTEGER] AS salaryK#4]]
+     *     |       | \_EsRelation[test][_meta_field{f}#32, emp_no{f}#26, first_name{f}#27, ..]
+     *     |       \_Aggregate[[salaryK{r}#4],[COUNT(*[KEYWORD],true[BOOLEAN]) AS count#6, salaryK{r}#4]]
+     *     |         \_StubRelation[[_meta_field{f}#32, emp_no{f}#26, first_name{f}#27, gender{f}#28, hire_date{f}#33, job{f}#34,
+     *                      job.raw{f}#35, languages{f}#29, last_name{f}#30, long_noidx{f}#36, salary{f}#31, salaryK{r}#4]]
+     *     \_Aggregate[[salaryK{r}#4],[MIN($$MV_COUNT(langua>$MIN$0{r$}#37,true[BOOLEAN]) AS min#19, salaryK{r}#4]]
+     *       \_Eval[[MVCOUNT(languages{f}#29) AS $$MV_COUNT(langua>$MIN$0#37]]
+     *         \_StubRelation[[_meta_field{f}#32, emp_no{f}#26, first_name{f}#27, gender{f}#28, hire_date{f}#33, job{f}#34, job.raw{f}#35,
+     *              languages{f}#29, last_name{f}#30, long_noidx{f}#36, salary{f}#31, count{r}#6, salaryK{r}#4, sum{r}#13,
+     *              hire_date_string{r}#10, date{r}#15, $$MV_COUNT(langua>$MIN$0{r$}#37]]
      */
     public void testTripleInlineStatsMultipleAssignmentsGetsPrunedPartially() {
         // TODO: reenable 1st sort, pull the 2nd further up when #132417 is in
@@ -6121,7 +6125,7 @@ public void testTripleInlineStatsMultipleAssignmentsGetsPrunedPartially() {
             "salary"
         );
 
-        var project = as(plan, EsqlProject.class);
+        var project = as(plan, Project.class);
         assertThat(Expressions.names(project.projections()), is(List.of("emp_no", "salaryK", "count", "min")));
         var topN = as(project.child(), TopN.class);
         var outerinline = as(topN.child(), InlineJoin.class);
@@ -6130,7 +6134,8 @@ public void testTripleInlineStatsMultipleAssignmentsGetsPrunedPartially() {
         expectedOutterOutput.addAll(List.of("count", "hire_date_string", "date", "min", "salaryK"));
         assertThat(Expressions.names(outerinline.output()), is(expectedOutterOutput));
         // outer left
-        var dissect = as(outerinline.left(), Dissect.class);
+        var outerProject = as(outerinline.left(), Project.class);
+        var dissect = as(outerProject.child(), Dissect.class);
         var eval = as(dissect.child(), Eval.class);
         var innerinline = as(eval.child(), InlineJoin.class);
         var expectedInnerOutput = new ArrayList<>(employeesFields);
@@ -6176,7 +6181,7 @@ public void testTripleInlineStatsMultipleAssignmentsGetsPrunedEntirely() {
         }
         var plan = optimizedPlan(query);
 
-        var project = as(plan, EsqlProject.class);
+        var project = as(plan, Project.class);
         assertThat(Expressions.names(project.projections()), is(List.of("emp_no")));
         var topN = as(project.child(), TopN.class);
         var dissect = as(topN.child(), Dissect.class);
@@ -6425,7 +6430,7 @@ public void testInlineStatsWithRow() {
         }
         var plan = optimizedPlan(query);
 
-        var esqlProject = as(plan, EsqlProject.class);
+        var esqlProject = as(plan, Project.class);
         assertThat(Expressions.names(esqlProject.projections()), is(List.of("emp_no")));
         var limit = asLimit(esqlProject.child(), 1000, false);
         var localRelation = as(limit.child(), LocalRelation.class);
@@ -6514,6 +6519,30 @@ public void testInlineStatsWithLimitAndAgg() {
         var stub = as(agg.child(), StubRelation.class);
     }
 
+    /*
+     * EsqlProject[[c{r}#7, b{r}#5, a{r}#14]]
+     * \_Eval[[[KEYWORD] AS a#14]]
+     *   \_Limit[1000[INTEGER],false]
+     *     \_LocalRelation[[a{r}#3, b{r}#5, c{r}#7],Page{blocks=[ConstantNullBlock[positions=1], IntVectorBlock[vector=ConstantIntVector[p
+     *          ositions=1, value=0]], BytesRefVectorBlock[vector=ConstantBytesRefVector[positions=1, value=[]]]]}]
+     */
+    public void testInlineStatsWithShadowedOutput() {
+        var query = """
+            ROW a = null, b = 0, c = ""
+            | INLINE STATS a = MAX(c) BY b
+            | EVAL a = c
+            """;
+        if (releaseBuildForInlineStats(query)) {
+            return;
+        }
+        var plan = optimizedPlan(query);
+        var project = as(plan, Project.class);
+        assertThat(Expressions.names(project.projections()), is(List.of("c", "b", "a")));
+        var eval = as(project.child(), Eval.class);
+        var limit = asLimit(eval.child(), 1000, false);
+        var localRelation = as(limit.child(), LocalRelation.class);
+    }
+
     /**
      * Expects
      *