Fix field resolution when one FORK branch is not bounded by a aggregation or project

Author: ioanatia

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

@@ -583,6 +583,29 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
             }
             projectAll.set(true);
         });
+
+        if (projectAll.get()) {
+            return result.withFieldNames(IndexResolver.ALL_FIELDS);
+        }
+
+        Holder<Boolean> projectAfterFork = new Holder<>(false);
+        Holder<Boolean> hasFork = new Holder<>(false);
+
+        parsed.forEachDown(plan -> {
+            if (hasFork.get() == false && (plan instanceof Project || plan instanceof Aggregate)) {
+                projectAfterFork.set(true);
+            }
+
+            if (plan instanceof Fork fork && projectAfterFork.get() == false) {
+                hasFork.set(true);
+                fork.children().forEach(child -> {
+                    if (child.anyMatch(p -> p instanceof Project || p instanceof Aggregate) == false) {
+                        projectAll.set(true);
+                    }
+                });
+            }
+        });
+
         if (projectAll.get()) {
             return result.withFieldNames(IndexResolver.ALL_FIELDS);
         }

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

@@ -1796,7 +1796,7 @@ public void testDropWildcardFields_WithLookupJoin() {
         );
     }
 
-    public void testForkFieldsWithKeep() {
+    public void testForkFieldsWithKeepAfterFork() {
         assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
 
         assertFieldNames("""
@@ -1810,6 +1810,20 @@ public void testForkFieldsWithKeep() {
             """, Set.of("a", "a.*", "c", "c.*", "d", "d.*", "e", "e.*", "x", "x.*", "y", "y.*"));
     }
 
+    public void testForkFieldsWithKeepBeforeFork() {
+        assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
+
+        assertFieldNames("""
+            FROM test
+            | KEEP a, b, c, d, x
+            | 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.*"));
+    }
+
     public void testForkFieldsWithNoProjection() {
         assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
 
@@ -1833,23 +1847,62 @@ public void testForkFieldsWithStatsInOneBranch() {
             | FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
                    (STATS x = count(*), y=min(z))
             | WHERE x > y
-            """, Set.of("a", "a.*", "c", "c.*", "z", "z.*"));
+            """, ALL_FIELDS);
     }
 
     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.*"
+            )
+        );
+    }
+
+    public void testForkWithStatsInAllBranches() {
+        assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
+
         assertFieldNames("""
             FROM test
-            | ENRICH enrich_policy ON abc
+            | WHERE a > 2000
             | EVAL b = a + 100
-            | LOOKUP JOIN my_lookup_index ON def
-            | FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
+            | FORK (WHERE c > 1 AND a < 10000 | STATS m = count(*))
+                   (EVAL z = a * b | STATS m = max(z))
                    (STATS x = count(*), y=min(z))
-            | LOOKUP JOIN my_lookup_index ON xyz
             | WHERE x > y
-            """, Set.of("a", "a.*", "abc", "abc.*", "c", "c.*", "def", "def.*", "x", "x.*", "xyz", "xyz.*", "y", "y.*", "z", "z.*"));
+            """, Set.of("a", "a.*", "b", "b.*", "c", "c.*", "z", "z.*"));
     }
 
     private Set<String> fieldNames(String query, Set<String> enrichPolicyMatchFields) {