Initial

Author: svilen-mihaylov-elastic

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

@@ -836,11 +836,6 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
             return result.withFieldNames(IndexResolver.ALL_FIELDS);
         }
 
-        // 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> projectAll = new Holder<>(false);
         parsed.forEachExpressionDown(UnresolvedStar.class, us -> {// explicit "*" fields selection
             if (projectAll.get()) {
@@ -853,6 +848,32 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
             return result.withFieldNames(IndexResolver.ALL_FIELDS);
         }
 
+        Holder<Boolean> projectAfterFork = new Holder<>(false);
+        Holder<Boolean> hasFork = new Holder<>(false);
+
+        parsed.forEachDown(plan -> {
+            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);
+                    }
+                });
+            }
+        });
+
+        if (projectAll.get()) {
+            return result.withFieldNames(IndexResolver.ALL_FIELDS);
+        }
+
         var referencesBuilder = AttributeSet.builder();
         // "keep" and "drop" attributes are special whenever a wildcard is used in their name, as the wildcard can cover some
         // attributes ("lookup join" generated columns among others); steps like removal of Aliases should ignore fields matching the

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

@@ -2080,7 +2080,7 @@ public void testForkFieldsWithKeepAfterFork() {
                    (WHERE d > 1000 AND e == "aaa" | EVAL c = a + 200)
             | WHERE x > y
             | KEEP a, b, c, d, x
-            """, ALL_FIELDS);
+            """, Set.of("a", "x", "y", "d", "e", "e.*", "d.*", "y.*", "x.*", "a.*"));
     }
 
     public void testForkFieldsWithKeepBeforeFork() {
@@ -2092,7 +2092,7 @@ public void testForkFieldsWithKeepBeforeFork() {
             | FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
                    (WHERE d > 1000 AND e == "aaa" | EVAL c = a + 200)
             | WHERE x > y
-            """, ALL_FIELDS);
+            """, Set.of("x", "y", "a", "d", "e", "b", "c", "e.*", "d.*", "y.*", "x.*", "a.*", "c.*", "b.*"));
     }
 
     public void testForkFieldsWithNoProjection() {
@@ -2118,17 +2118,41 @@ public void testForkFieldsWithStatsInOneBranch() {
     }
 
     public void testForkFieldsWithEnrichAndLookupJoins() {
-        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);
+        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.*"
+            )
+        );
     }
 
     public void testForkWithStatsInAllBranches() {
@@ -2140,7 +2164,7 @@ public void testForkWithStatsInAllBranches() {
                    (EVAL z = a * b | STATS m = max(z))
                    (STATS x = count(*), y=min(z))
             | WHERE x > y
-            """, ALL_FIELDS);
+            """, Set.of("c", "a", "z", "z.*", "a.*", "c.*"));
     }
 
     public void testForkWithStatsAndWhere() {