ES|QL: Fix Fork field reference tracking

docs/changelog/131723.yaml

@@ -0,0 +1,5 @@
+pr: 131723
+summary: Tests for FORK's evaluation of field names used in `field_caps` resolve calls
+area: Search
+type: bug
+issues: []

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/tree/Node.java

@@ -65,13 +65,28 @@ public List<T> children() {
     }
 
     @SuppressWarnings("unchecked")
-    public void forEachDown(Consumer<? super T> action) {
-        action.accept((T) this);
+    public boolean forEachDownMayReturnEarly(Function<? super T, Boolean> action) {
+        if (action.apply((T) this) == false) {
+            // Early return.
+            return false;
+        }
         // please do not refactor it to a for-each loop to avoid
         // allocating iterator that performs concurrent modification checks and extra stack frames
         for (int c = 0, size = children.size(); c < size; c++) {
-            children.get(c).forEachDown(action);
+            if (children.get(c).forEachDownMayReturnEarly(action) == false) {
+                return false;
+            }
         }
+        return true;
+    }
+
+    @SuppressWarnings("unchecked")
+    public void forEachDown(Consumer<? super T> action) {
+        forEachDownMayReturnEarly(p -> {
+            action.accept(p);
+            // No early return.
+            return true;
+        });
     }
 
     @SuppressWarnings("unchecked")

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

@@ -49,6 +49,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.xpack.esql.core.util.StringUtils.WILDCARD;
@@ -76,11 +77,6 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
             return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
         }
 
-        // TODO: Improve field resolution for FORK - right now we request all fields
-        if (parsed.anyMatch(p -> p instanceof Fork)) {
-            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
-        }
-
         Holder<Boolean> projectAll = new Holder<>(false);
         parsed.forEachExpressionDown(UnresolvedStar.class, us -> {// explicit "*" fields selection
             if (projectAll.get()) {
@@ -93,7 +89,7 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
             return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
         }
 
-        var referencesBuilder = AttributeSet.builder();
+        var referencesBuilder = new Holder<>(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
         // wildcards.
@@ -110,19 +106,42 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
         // lookup indices where we request "*" because we may require all their fields
         Set<String> wildcardJoinIndices = new java.util.HashSet<>();
 
-        boolean[] canRemoveAliases = new boolean[] { true };
+        var canRemoveAliases = new Holder<>(true);
+        var needsAllFields = new Holder<>(false);
 
-        parsed.forEachDown(p -> {// go over each plan top-down
-            if (p instanceof RegexExtract re) { // for Grok and Dissect
+        var processingLambda = new Holder<Function<LogicalPlan, Boolean>>();
+        processingLambda.set((LogicalPlan p) -> {// go over each plan top-down
+            if (p instanceof Fork fork) {
+                // Early return from forEachDown. We will iterate over the children manually.
+                var forkRefsResult = AttributeSet.builder();
+                forkRefsResult.addAll(referencesBuilder.get());
+
+                for (var child : fork.children()) {
+                    referencesBuilder.set(AttributeSet.builder());
+                    var return_result = child.forEachDownMayReturnEarly(processingLambda.get());
+                    // No nested Forks for now...
+                    assert return_result;
+                    if (referencesBuilder.get().isEmpty()) {
+                        needsAllFields.set(true);
+                        // Early return.
+                        return false;
+                    }
+                    forkRefsResult.addAll(referencesBuilder.get());
+                }
+
+                forkRefsResult.removeIf(attr -> attr.name().equals(Fork.FORK_FIELD));
+                referencesBuilder.set(forkRefsResult);
+                return false;
+            } else if (p instanceof RegexExtract re) { // for Grok and Dissect
                 // keep the inputs needed by Grok/Dissect
-                referencesBuilder.addAll(re.input().references());
+                referencesBuilder.get().addAll(re.input().references());
             } else if (p instanceof Enrich enrich) {
                 AttributeSet enrichFieldRefs = Expressions.references(enrich.enrichFields());
                 AttributeSet.Builder enrichRefs = enrichFieldRefs.combine(enrich.matchField().references()).asBuilder();
                 // Enrich adds an EmptyAttribute if no match field is specified
                 // The exact name of the field will be added later as part of enrichPolicyMatchFields Set
                 enrichRefs.removeIf(attr -> attr instanceof EmptyAttribute);
-                referencesBuilder.addAll(enrichRefs);
+                referencesBuilder.get().addAll(enrichRefs);
             } else if (p instanceof LookupJoin join) {
                 if (join.config().type() instanceof JoinTypes.UsingJoinType usingJoinType) {
                     joinRefs.addAll(usingJoinType.columns());
@@ -135,15 +154,15 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
                     joinRefs.addAll(keepRefs);
                 }
             } else {
-                referencesBuilder.addAll(p.references());
+                referencesBuilder.get().addAll(p.references());
                 if (p instanceof UnresolvedRelation ur && ur.indexMode() == IndexMode.TIME_SERIES) {
                     // METRICS aggs generally rely on @timestamp without the user having to mention it.
-                    referencesBuilder.add(new UnresolvedAttribute(ur.source(), MetadataAttribute.TIMESTAMP_FIELD));
+                    referencesBuilder.get().add(new UnresolvedAttribute(ur.source(), MetadataAttribute.TIMESTAMP_FIELD));
                 }
                 // special handling for UnresolvedPattern (which is not an UnresolvedAttribute)
                 p.forEachExpression(UnresolvedNamePattern.class, up -> {
                     var ua = new UnresolvedAttribute(up.source(), up.name());
-                    referencesBuilder.add(ua);
+                    referencesBuilder.get().add(ua);
                     if (p instanceof Keep) {
                         keepRefs.add(ua);
                     } else if (p instanceof Drop) {
@@ -168,10 +187,10 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
             //
             // and ips_policy enriches the results with the same name ip field),
             // these aliases should be kept in the list of fields.
-            if (canRemoveAliases[0] && p.anyMatch(FieldNameUtils::couldOverrideAliases)) {
-                canRemoveAliases[0] = false;
+            if (canRemoveAliases.get() && p.anyMatch(FieldNameUtils::couldOverrideAliases)) {
+                canRemoveAliases.set(false);
             }
-            if (canRemoveAliases[0]) {
+            if (canRemoveAliases.get()) {
                 // remove any already discovered UnresolvedAttributes that are in fact aliases defined later down in the tree
                 // for example "from test | eval x = salary | stats max = max(x) by gender"
                 // remove the UnresolvedAttribute "x", since that is an Alias defined in "eval"
@@ -187,21 +206,28 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
                     if (fieldNames.contains(ne.name())) {
                         return;
                     }
-                    referencesBuilder.removeIf(
-                        attr -> matchByName(attr, ne.name(), keepRefs.contains(attr) || dropWildcardRefs.contains(attr))
-                    );
+                    referencesBuilder.get()
+                        .removeIf(attr -> matchByName(attr, ne.name(), keepRefs.contains(attr) || dropWildcardRefs.contains(attr)));
                 });
             }
+
+            // No early return.
+            return true;
         });
+        parsed.forEachDownMayReturnEarly(processingLambda.get());
+
+        if (needsAllFields.get()) {
+            return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
+        }
 
         // Add JOIN ON column references afterward to avoid Alias removal
-        referencesBuilder.addAll(joinRefs);
+        referencesBuilder.get().addAll(joinRefs);
         // If any JOIN commands need wildcard field-caps calls, persist the index names
 
         // remove valid metadata attributes because they will be filtered out by the IndexResolver anyway
         // otherwise, in some edge cases, we will fail to ask for "*" (all fields) instead
-        referencesBuilder.removeIf(a -> a instanceof MetadataAttribute || MetadataAttribute.isSupported(a.name()));
-        Set<String> fieldNames = referencesBuilder.build().names();
+        referencesBuilder.get().removeIf(a -> a instanceof MetadataAttribute || MetadataAttribute.isSupported(a.name()));
+        Set<String> fieldNames = referencesBuilder.get().build().names();
 
         if (fieldNames.isEmpty() && enrichPolicyMatchFields.isEmpty()) {
             // there cannot be an empty list of fields, we'll ask the simplest and lightest one instead: _index

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

@@ -59,6 +59,13 @@ public void testSimple1() {
         );
     }
 
+    public void testSimple2() {
+        assertFieldNames("""
+            FROM employees
+            | WHERE emp_no == "2"
+            """, IndexResolver.ALL_FIELDS);
+    }
+
     public void testDirectFilter() {
         assertFieldNames(
             "from employees | sort emp_no | where still_hired | keep emp_no | limit 3",
@@ -2074,7 +2081,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", "c", "d", "e", "e.*", "d.*", "y.*", "x.*", "a.*", "c.*"));
     }
 
     public void testForkFieldsWithKeepBeforeFork() {
@@ -2086,7 +2093,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() {
@@ -2108,21 +2115,25 @@ public void testForkFieldsWithStatsInOneBranch() {
             | FORK (WHERE c > 1 AND a < 10000 | EVAL d = a + 500)
                    (STATS x = count(*), y=min(z))
             | WHERE x > y
-            """, ALL_FIELDS);
+            """, Set.of("x", "y", "a", "c", "z", "y.*", "x.*", "z.*", "a.*", "c.*"));
     }
 
     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", "a", "c", "abc", "b", "def", "z", "xyz", "def.*", "y.*", "x.*", "xyz.*", "z.*", "abc.*", "a.*", "c.*", "b.*"),
+            Set.of("my_lookup_index")
+        );
     }
 
     public void testForkWithStatsInAllBranches() {
@@ -2134,11 +2145,86 @@ public void testForkWithStatsInAllBranches() {
                    (EVAL z = a * b | STATS m = max(z))
                    (STATS x = count(*), y=min(z))
             | WHERE x > y
-            """, ALL_FIELDS);
+            """, Set.of("x", "y", "c", "a", "z", "y.*", "x.*", "z.*", "a.*", "c.*"));
+    }
+
+    public void testForkWithStatsInAllBranches1() {
+        assertFieldNames("""
+            FROM employees
+            | FORK
+                   ( STATS x = min(last_name))
+                   ( EVAL last_name = first_name  | STATS y = max(last_name))
+            """, Set.of("first_name", "last_name", "first_name.*", "last_name.*"));
+    }
+
+    public void testForkWithStatsInAllBranches2() {
+        assertFieldNames("""
+            FROM employees
+            | FORK
+                   ( EVAL last_name = first_name  | STATS y = VALUES(last_name))
+                   ( STATS x = VALUES(last_name))
+            """, Set.of("first_name", "last_name", "first_name.*", "last_name.*"));
     }
 
     public void testForkWithStatsAndWhere() {
-        assertFieldNames(" FROM employees | FORK ( WHERE true | stats min(salary) by gender) ( WHERE true | LIMIT 3 )", ALL_FIELDS);
+        assertFieldNames(
+            " FROM employees | FORK ( WHERE true | stats min(salary) by gender) ( WHERE true | LIMIT 3 )",
+            IndexResolver.ALL_FIELDS
+        );
+    }
+
+    public void testNullString() {
+        assertFieldNames("""
+             FROM sample_data
+            | EVAL x = null::string
+            | STATS COUNT() BY category=CATEGORIZE(x)
+            | SORT category""", Set.of("_index"));
+    }
+
+    public void testNullStringWithFork() {
+        assertFieldNames("""
+             FROM sample_data
+            | EVAL x = null::string
+            | STATS COUNT() BY category=CATEGORIZE(x)
+            | SORT category
+            | FORK (WHERE true) (WHERE true) | WHERE _fork == "fork1" | DROP _fork""", IndexResolver.ALL_FIELDS);
+    }
+
+    public void testSingleFork() {
+        assertFieldNames("""
+             FROM employees
+            | FORK
+               ( STATS x = count(*))
+               ( WHERE emp_no == "2" )
+            | SORT _fork""", IndexResolver.ALL_FIELDS);
+    }
+
+    public void testForkRefs1() {
+        assertFieldNames("""
+            FROM employees
+            | KEEP first_name, last_name
+            | FORK
+               ( EVAL x = first_name)
+               ( EVAL x = last_name)
+            """, Set.of("first_name", "last_name", "last_name.*", "first_name.*"));
+    }
+
+    public void testForkRefs2() {
+        assertFieldNames("""
+            FROM employees
+            | FORK
+               ( KEEP first_name | EVAL x = first_name)
+               ( KEEP last_name | EVAL x = last_name)
+            """, Set.of("first_name", "last_name", "last_name.*", "first_name.*"));
+    }
+
+    public void testForkRefs3() {
+        assertFieldNames("""
+            FROM employees
+            | FORK
+               ( KEEP first_name | EVAL last_name = first_name)
+               ( KEEP first_name | EVAL x = first_name)
+            """, Set.of("first_name", "first_name.*"));
     }
 
     private void assertFieldNames(String query, Set<String> expected) {