ESQL: Fix inconsistent column order in MV_EXPAND (elastic#129745) (elastic#131413)

The new attribute generated by MV_EXPAND should remain in the original position. The projection added by ProjectAwayColumns does not respect the original order of attributes.

Make ProjectAwayColumns respect the order of attributes to fix this.

Co-authored-by: kanoshiou <[email protected]>

Author: alex-spies

docs/changelog/129745.yaml

@@ -0,0 +1,6 @@
+pr: 129745
+summary: "ESQL: Fix `mv_expand` inconsistent column order"
+area: ES|QL
+type: bug
+issues:
+  - 129000

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/AttributeSet.java

@@ -261,5 +261,9 @@ public boolean isEmpty() {
         public AttributeSet build() {
             return new AttributeSet(mapBuilder.build());
         }
+
+        public void clear() {
+            mapBuilder.keySet().clear();
+        }
     }
 }

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

@@ -53,7 +53,6 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
         "Data too large", // Circuit breaker exceptions eg. https://github.com/elastic/elasticsearch/issues/130072
 
         // Awaiting fixes for correctness
-        "Expecting the following columns \\[.*\\], got", // https://github.com/elastic/elasticsearch/issues/129000
         "Expecting at most \\[.*\\] columns, got \\[.*\\]" // https://github.com/elastic/elasticsearch/issues/129561
     );
 

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

@@ -419,3 +419,65 @@ emp_no:integer | job_positions:keyword
 10001          | Accountant
 10001          | Senior Python Developer          
 ;
+
+testMvExpandInconsistentColumnOrder1
+required_capability: fix_mv_expand_inconsistent_column_order
+from message_types
+| eval foo_1 = 1, foo_2 = 2 
+| sort message
+| mv_expand foo_1
+;
+
+message:keyword         | type:keyword | foo_1:integer | foo_2:integer
+Connected to 10.1.0.1   | Success      | 1             | 2
+Connected to 10.1.0.2   | Success      | 1             | 2
+Connected to 10.1.0.3   | Success      | 1             | 2
+Connection error        | Error        | 1             | 2
+Development environment | Development  | 1             | 2
+Disconnected            | Disconnected | 1             | 2
+Production environment  | Production   | 1             | 2
+;
+
+testMvExpandInconsistentColumnOrder2
+required_capability: fix_mv_expand_inconsistent_column_order
+from message_types
+| eval foo_1 = [1, 3], foo_2 = 2
+| sort message
+| mv_expand foo_1
+;
+
+message:keyword         | type:keyword | foo_1:integer | foo_2:integer
+Connected to 10.1.0.1   | Success      | 1             | 2
+Connected to 10.1.0.1   | Success      | 3             | 2
+Connected to 10.1.0.2   | Success      | 1             | 2
+Connected to 10.1.0.2   | Success      | 3             | 2
+Connected to 10.1.0.3   | Success      | 1             | 2
+Connected to 10.1.0.3   | Success      | 3             | 2
+Connection error        | Error        | 1             | 2
+Connection error        | Error        | 3             | 2
+Development environment | Development  | 1             | 2
+Development environment | Development  | 3             | 2
+Disconnected            | Disconnected | 1             | 2
+Disconnected            | Disconnected | 3             | 2
+Production environment  | Production   | 1             | 2
+Production environment  | Production   | 3             | 2
+;
+
+testMvExpandInconsistentColumnOrder3
+required_capability: fix_mv_expand_inconsistent_column_order
+from message_types
+| sort type
+| eval language_code = 1, `language_name` = false, message = true, foo_3 = 1, foo_2 = null
+| eval foo_3 = "1", `foo_3` = -1, foo_1 = 1, `language_code` = null, `foo_2` = "1"
+| mv_expand foo_1
+| limit 5
+;
+
+type:keyword | language_name:boolean | message:boolean | foo_3:integer | foo_1:integer | language_code:null | foo_2:keyword
+Development  | false                 | true            | -1            | 1             | null               | 1
+Disconnected | false                 | true            | -1            | 1             | null               | 1
+Error        | false                 | true            | -1            | 1             | null               | 1
+Production   | false                 | true            | -1            | 1             | null               | 1
+Success      | false                 | true            | -1            | 1             | null               | 1
+;
+

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

@@ -1227,6 +1227,12 @@ public enum Cap {
          */
         NO_PLAIN_STRINGS_IN_LITERALS,
 
+        /**
+         * Support for the mv_expand target attribute should be retained in its original position.
+         * see <a href="https://github.com/elastic/elasticsearch/issues/129000"> ES|QL: inconsistent column order #129000 </a>
+         */
+        FIX_MV_EXPAND_INCONSISTENT_COLUMN_ORDER,
+
         /**
          * (Re)Added EXPLAIN command
          */

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/ProjectAwayColumns.java

@@ -76,7 +76,19 @@ public PhysicalPlan apply(PhysicalPlan plan) {
 
                     // no need for projection when dealing with aggs
                     if (logicalFragment instanceof Aggregate == false) {
-                        List<Attribute> output = new ArrayList<>(requiredAttrBuilder.build());
+                        // we should respect the order of the attributes
+                        List<Attribute> output = new ArrayList<>();
+                        for (Attribute attribute : logicalFragment.output()) {
+                            if (requiredAttrBuilder.contains(attribute)) {
+                                output.add(attribute);
+                                requiredAttrBuilder.remove(attribute);
+                            }
+                        }
+                        // requiredAttrBuilder should be empty unless the plan is inconsistent due to a bug.
+                        // This can happen in case of remote ENRICH, see https://github.com/elastic/elasticsearch/issues/118531
+                        // TODO: stop adding the remaining required attributes once remote ENRICH is fixed.
+                        output.addAll(requiredAttrBuilder.build());
+
                         // if all the fields are filtered out, it's only the count that matters
                         // however until a proper fix (see https://github.com/elastic/elasticsearch/issues/98703)
                         // add a synthetic field (so it doesn't clash with the user defined one) to return a constant