Add test and apply refactor suggestion

Author: ioanatia

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

@@ -571,7 +571,7 @@ private void resolveInferences(
     }
 
     static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicyMatchFields, PreAnalysisResult result) {
-        if (false == parsed.anyMatch(plan -> plan instanceof Aggregate || plan instanceof Project)) {
+        if (false == parsed.anyMatch(EsqlSession::shouldCollectReferencedFields)) {
             // no explicit columns selection, for example "from employees"
             return result.withFieldNames(IndexResolver.ALL_FIELDS);
         }
@@ -592,14 +592,18 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
         Holder<Boolean> hasFork = new Holder<>(false);
 
         parsed.forEachDown(plan -> {
-            if (hasFork.get() == false && (plan instanceof Project || plan instanceof Aggregate)) {
+            if (projectAll.get()) {
+                return;
+            }
+
+            if (hasFork.get() == false && shouldCollectReferencedFields(plan)) {
                 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) {
+                    if (child.anyMatch(EsqlSession::shouldCollectReferencedFields) == false) {
                         projectAll.set(true);
                     }
                 });
@@ -722,6 +726,13 @@ static PreAnalysisResult fieldNames(LogicalPlan parsed, Set<String> enrichPolicy
         }
     }
 
+    /**
+     * Indicates whether the given plan gives an exact list of fields that we need to collect from field_caps.
+     */
+    private static boolean shouldCollectReferencedFields(LogicalPlan plan) {
+        return plan instanceof Aggregate || plan instanceof Project;
+    }
+
     /**
      * Could a plan "accidentally" override aliases?
      * Examples are JOIN and ENRICH, that _could_ produce fields with the same

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

@@ -1905,6 +1905,12 @@ public void testForkWithStatsInAllBranches() {
             """, Set.of("a", "a.*", "b", "b.*", "c", "c.*", "z", "z.*"));
     }
 
+    public void testForkWithStatsAndWhere() {
+        assumeTrue("FORK available as snapshot only", EsqlCapabilities.Cap.FORK.isEnabled());
+
+        assertFieldNames(" FROM employees | FORK ( WHERE true | stats min(salary) by gender) ( WHERE true | LIMIT 3 )", ALL_FIELDS);
+    }
+
     private Set<String> fieldNames(String query, Set<String> enrichPolicyMatchFields) {
         var preAnalysisResult = new EsqlSession.PreAnalysisResult(null);
         return EsqlSession.fieldNames(parser.createStatement(query), enrichPolicyMatchFields, preAnalysisResult).fieldNames();