ESQL: Fix resolution of MV_EXPAND after KEEP * (#103339)

docs/changelog/103339.yaml

@@ -0,0 +1,6 @@
+pr: 103339
+summary: "ESQL: Fix resolution of MV_EXPAND after KEEP *"
+area: ES|QL
+type: bug
+issues:
+ - 103331

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

@@ -316,3 +316,11 @@ a:keyword | e:keyword
 a         | a
 ;
 
+
+//see https://github.com/elastic/elasticsearch/issues/103331
+keepStarMvExpand#[skip:-8.12.99]
+from employees | where emp_no == 10001 | keep * |  mv_expand first_name;
+
+avg_worked_seconds:long | birth_date:date          | emp_no:integer | first_name:keyword | gender:keyword | height:double | height.float:double | height.half_float:double | height.scaled_float:double | hire_date:date           | is_rehired:boolean | job_positions:keyword                 | languages:integer | languages.byte:integer | languages.long:long | languages.short:integer | last_name:keyword | salary:integer | salary_change:double | salary_change.int:integer | salary_change.keyword:keyword | salary_change.long:long | still_hired:boolean  
+268728049               | 1953-09-02T00:00:00.000Z | 10001          | Georgi             | M              | 2.03          | 2.0299999713897705  | 2.029296875              | 2.0300000000000002         | 1986-06-26T00:00:00.000Z | [false, true]      | [Accountant, Senior Python Developer] | 2                 | 2                      | 2                   | 2                       | Facello           | 57305          | 1.19                 | 1                         | 1.19                          | 1                       | true           
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/MvExpand.java

@@ -22,13 +22,12 @@ public class MvExpand extends UnaryPlan {
     private final NamedExpression target;
     private final Attribute expanded;
 
-    private final List<Attribute> output;
+    private List<Attribute> output;
 
     public MvExpand(Source source, LogicalPlan child, NamedExpression target, Attribute expanded) {
         super(source, child);
         this.target = target;
         this.expanded = expanded;
-        this.output = calculateOutput(child.output(), target, expanded);
     }
 
     public static List<Attribute> calculateOutput(List<Attribute> input, NamedExpression target, Attribute expanded) {
@@ -63,6 +62,9 @@ public class MvExpand extends UnaryPlan {
 
     @Override
     public List<Attribute> output() {
+        if (output == null) {
+            output = calculateOutput(child().output(), target, expanded);
+        }
         return output;
     }
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/parser/StatementParserTests.java

@@ -26,6 +26,7 @@ import org.elasticsearch.xpack.esql.plan.logical.Grok;
 import org.elasticsearch.xpack.esql.plan.logical.InlineStats;
 import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
 import org.elasticsearch.xpack.esql.plan.logical.Row;
+import org.elasticsearch.xpack.ql.capabilities.UnresolvedException;
 import org.elasticsearch.xpack.ql.expression.Alias;
 import org.elasticsearch.xpack.ql.expression.EmptyAttribute;
 import org.elasticsearch.xpack.ql.expression.Expressions;
@@ -710,6 +711,17 @@ public class StatementParserTests extends ESTestCase {
         assertThat(expand.target(), equalTo(attribute("a")));
     }
 
+    // see https://github.com/elastic/elasticsearch/issues/103331
+    public void testKeepStarMvExpand() {
+        try {
+            String query = "from test | keep * | mv_expand a";
+            var plan = statement(query);
+        } catch (UnresolvedException e) {
+            fail(e, "Regression: https://github.com/elastic/elasticsearch/issues/103331");
+        }
+
+    }
+
     public void testUsageOfProject() {
         processingCommand("project a");
         assertWarnings("PROJECT command is no longer supported, please use KEEP instead");