ES|QL: Improve field resolution for FORK (#128501)

Author: ioanatia

muted-tests.yml

@@ -453,12 +453,6 @@ tests:
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test121CanUseStackLoggingConfig
   issue: https://github.com/elastic/elasticsearch/issues/128165
-- class: org.elasticsearch.xpack.esql.session.IndexResolverFieldNamesTests
-  method: testForkFieldsWithKeepAfterFork
-  issue: https://github.com/elastic/elasticsearch/issues/128271
-- class: org.elasticsearch.xpack.esql.session.IndexResolverFieldNamesTests
-  method: testForkWithStatsInAllBranches
-  issue: https://github.com/elastic/elasticsearch/issues/128272
 - class: org.elasticsearch.packaging.test.DockerTests
   method: test080ConfigurePasswordThroughEnvironmentVariableFile
   issue: https://github.com/elastic/elasticsearch/issues/128075

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

@@ -584,38 +584,17 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
             return result.withFieldNames(IndexResolver.ALL_FIELDS);
         }
 
-        Holder<Boolean> projectAll = new Holder<>(false);
-        parsed.forEachExpressionDown(UnresolvedStar.class, us -> {// explicit "*" fields selection
-            if (projectAll.get()) {
-                return;
-            }
-            projectAll.set(true);
-        });
-
-        if (projectAll.get()) {
+        // TODO: Improve field resolution for FORK - right now we request all fields
+        if (parsed.anyMatch(p -> p instanceof Fork)) {
             return result.withFieldNames(IndexResolver.ALL_FIELDS);
         }
 
-        Holder<Boolean> projectAfterFork = new Holder<>(false);
-        Holder<Boolean> hasFork = new Holder<>(false);
-
-        parsed.forEachDown(plan -> {
+        Holder<Boolean> projectAll = new Holder<>(false);
+        parsed.forEachExpressionDown(UnresolvedStar.class, us -> {// explicit "*" fields selection
             if (projectAll.get()) {
                 return;
             }
-
-            if (hasFork.get() == false && shouldCollectReferencedFields(plan, inlinestatsAggs)) {
-                projectAfterFork.set(true);
-            }
-
-            if (plan instanceof Fork fork && projectAfterFork.get() == false) {
-                hasFork.set(true);
-                fork.children().forEach(child -> {
-                    if (child.anyMatch(p -> shouldCollectReferencedFields(p, inlinestatsAggs)) == false) {
-                        projectAll.set(true);
-                    }
-                });
-            }
+            projectAll.set(true);
         });
 
         if (projectAll.get()) {

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

@@ -2009,21 +2009,21 @@ public void testForkFieldsWithKeepAfterFork() {
                    (WHERE d > 1000 AND e == "aaa" | EVAL c = a + 200)
             | WHERE x > y
             | KEEP a, b, c, d, x
-            """, Set.of("a", "a.*", "c", "c.*", "d", "d.*", "e", "e.*", "x", "x.*", "y", "y.*"));
+            """, ALL_FIELDS);
     }
 
     public void testForkFieldsWithKeepBeforeFork() {
         assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
 
         assertFieldNames("""
             FROM test
-            | KEEP a, b, c, d, x
+            | KEEP a, b, c, d, x, y
             | WHERE a > 2000
             | EVAL b = a + 100
             | FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
                    (WHERE d > 1000 AND e == "aaa" | EVAL c = a + 200)
             | WHERE x > y
-            """, Set.of("a", "a.*", "b", "b.*", "c", "c.*", "d", "d.*", "e", "e.*", "x", "x.*", "y", "y.*"));
+            """, ALL_FIELDS);
     }
 
     public void testForkFieldsWithNoProjection() {
@@ -2056,41 +2056,17 @@ public void testForkFieldsWithEnrichAndLookupJoins() {
         assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
         assumeTrue("LOOKUP JOIN available as snapshot only", EsqlCapabilities.Cap.JOIN_LOOKUP_V12.isEnabled());
 
-        assertFieldNames(
-            """
-                FROM test
-                | KEEP a, b, abc, def, z, xyz
-                | ENRICH enrich_policy ON abc
-                | EVAL b = a + 100
-                | LOOKUP JOIN my_lookup_index ON def
-                | FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
-                       (STATS x = count(*), y=min(z))
-                | LOOKUP JOIN my_lookup_index ON xyz
-                | WHERE x > y OR _fork == "fork1"
-                """,
-            Set.of(
-                "x",
-                "y",
-                "_fork",
-                "a",
-                "c",
-                "abc",
-                "b",
-                "def",
-                "z",
-                "xyz",
-                "def.*",
-                "_fork.*",
-                "y.*",
-                "x.*",
-                "xyz.*",
-                "z.*",
-                "abc.*",
-                "a.*",
-                "c.*",
-                "b.*"
-            )
-        );
+        assertFieldNames("""
+            FROM test
+            | KEEP a, b, abc, def, z, xyz
+            | ENRICH enrich_policy ON abc
+            | EVAL b = a + 100
+            | LOOKUP JOIN my_lookup_index ON def
+            | FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
+                   (STATS x = count(*), y=min(z))
+            | LOOKUP JOIN my_lookup_index ON xyz
+            | WHERE x > y OR _fork == "fork1"
+            """, ALL_FIELDS);
     }
 
     public void testForkWithStatsInAllBranches() {
@@ -2104,7 +2080,7 @@ public void testForkWithStatsInAllBranches() {
                    (EVAL z = a * b | STATS m = max(z))
                    (STATS x = count(*), y=min(z))
             | WHERE x > y
-            """, Set.of("a", "a.*", "b", "b.*", "c", "c.*", "z", "z.*"));
+            """, ALL_FIELDS);
     }
 
     public void testForkWithStatsAndWhere() {