Address feedback

Author: svilen-mihaylov-elastic

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

@@ -8,12 +8,14 @@
 
 import org.elasticsearch.common.io.stream.NamedWriteable;
 import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
+import org.elasticsearch.xpack.esql.core.util.Holder;
 
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -65,28 +67,33 @@ public List<T> children() {
     }
 
     @SuppressWarnings("unchecked")
-    public boolean forEachDownMayReturnEarly(Function<? super T, Boolean> action) {
-        if (action.apply((T) this) == false) {
+    void forEachDownMayReturnEarly(BiConsumer<? super T, Holder<Boolean>> action, Holder<Boolean> breakEarly) {
+        action.accept((T) this, breakEarly);
+        if (breakEarly.get()) {
             // Early return.
-            return false;
+            return;
         }
         // 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++) {
-            if (children.get(c).forEachDownMayReturnEarly(action) == false) {
-                return false;
+            children.get(c).forEachDownMayReturnEarly(action, breakEarly);
+            if (breakEarly.get()) {
+                // Early return.
+                return;
             }
         }
-        return true;
     }
 
-    @SuppressWarnings("unchecked")
+    public boolean forEachDownMayReturnEarly(BiConsumer<? super T, Holder<Boolean>> action) {
+        var breakEarly = new Holder<>(false);
+        forEachDownMayReturnEarly(action, breakEarly);
+        return breakEarly.get();
+    }
+
     public void forEachDown(Consumer<? super T> action) {
-        forEachDownMayReturnEarly(p -> {
-            action.accept(p);
-            // No early return.
-            return true;
-        });
+        boolean result = forEachDownMayReturnEarly((p, breakFlag) -> { action.accept(p); });
+        // We should be breaking early here...
+        assert result == false;
     }
 
     @SuppressWarnings("unchecked")

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

@@ -49,7 +49,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
-import java.util.function.Function;
+import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.xpack.esql.core.util.StringUtils.WILDCARD;
@@ -107,31 +107,36 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
         Set<String> wildcardJoinIndices = new java.util.HashSet<>();
 
         var canRemoveAliases = new Holder<>(true);
-        var needsAllFields = new Holder<>(false);
 
-        var processingLambda = new Holder<Function<LogicalPlan, Boolean>>();
-        processingLambda.set((LogicalPlan p) -> {// go over each plan top-down
+        var processingLambda = new Holder<BiConsumer<LogicalPlan, Holder<Boolean>>>();
+        processingLambda.set((LogicalPlan p, Holder<Boolean> breakEarly) -> {// go over each plan top-down
             if (p instanceof Fork fork) {
-                // Early return from forEachDown. We will iterate over the children manually.
+                // Early return from forEachDown. We will iterate over the children manually and end the recursion via forEachDown early.
                 var forkRefsResult = AttributeSet.builder();
                 forkRefsResult.addAll(referencesBuilder.get());
 
-                for (var child : fork.children()) {
+                for (var fork_child : fork.children()) {
                     referencesBuilder.set(AttributeSet.builder());
-                    var return_result = child.forEachDownMayReturnEarly(processingLambda.get());
-                    // No nested Forks for now...
-                    assert return_result;
+                    var nested_early_return = fork_child.forEachDownMayReturnEarly(processingLambda.get());
+                    // This assert is just for good measure. FORKs within FORKs is yet not supported.
+                    assert nested_early_return == false;
+
+                    // See below, no references, means we should return all fields (*).
                     if (referencesBuilder.get().isEmpty()) {
-                        needsAllFields.set(true);
-                        // Early return.
-                        return false;
+                        projectAll.set(true);
+                        // Return early, we'll be returning all references no matter what the remainder of the query is.
+                        breakEarly.set(true);
+                        return;
                     }
                     forkRefsResult.addAll(referencesBuilder.get());
                 }
 
                 forkRefsResult.removeIf(attr -> attr.name().equals(Fork.FORK_FIELD));
                 referencesBuilder.set(forkRefsResult);
-                return false;
+
+                // Return early, we've already explored all fork branches.
+                breakEarly.set(true);
+                return;
             } else if (p instanceof RegexExtract re) { // for Grok and Dissect
                 // keep the inputs needed by Grok/Dissect
                 referencesBuilder.get().addAll(re.input().references());
@@ -210,13 +215,10 @@ public static PreAnalysisResult resolveFieldNames(LogicalPlan parsed, EnrichReso
                         .removeIf(attr -> matchByName(attr, ne.name(), keepRefs.contains(attr) || dropWildcardRefs.contains(attr)));
                 });
             }
-
-            // No early return.
-            return true;
         });
         parsed.forEachDownMayReturnEarly(processingLambda.get());
 
-        if (needsAllFields.get()) {
+        if (projectAll.get()) {
             return new PreAnalysisResult(enrichResolution, IndexResolver.ALL_FIELDS, Set.of());
         }