ESQL: Reduce iteration complexity for plan traversal (elastic#123427) (elastic#123534)

(cherry picked from commit e4604a4)

Author: costin

docs/changelog/123427.yaml

@@ -0,0 +1,5 @@
+pr: 123427
+summary: Reduce iteration complexity for plan traversal
+area: ES|QL
+type: bug
+issues: []

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java

@@ -14,14 +14,14 @@
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.ElementType;
 import org.elasticsearch.core.Tuple;
+import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
-import org.elasticsearch.xpack.esql.core.tree.Node;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.util.Holder;
@@ -33,15 +33,13 @@
 import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
-import org.elasticsearch.xpack.esql.plan.logical.join.Join;
 import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EsSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.EstimatesRowSize;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSinkExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeSourceExec;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
-import org.elasticsearch.xpack.esql.plan.physical.LookupJoinExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.planner.mapper.LocalMapper;
 import org.elasticsearch.xpack.esql.planner.mapper.Mapper;
@@ -50,12 +48,10 @@
 import org.elasticsearch.xpack.esql.stats.SearchStats;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
 import java.util.function.Consumer;
-import java.util.function.Function;
 
 import static java.util.Arrays.asList;
 import static org.elasticsearch.index.mapper.MappedFieldType.FieldExtractPreference.DOC_VALUES;
@@ -110,7 +106,7 @@ public static Set<String> planConcreteIndices(PhysicalPlan plan) {
             return Set.of();
         }
         var indices = new LinkedHashSet<String>();
-        forEachFromRelation(plan, relation -> indices.addAll(relation.concreteIndices()));
+        forEachRelation(plan, relation -> indices.addAll(relation.concreteIndices()));
         return indices;
     }
 
@@ -122,42 +118,16 @@ public static String[] planOriginalIndices(PhysicalPlan plan) {
             return Strings.EMPTY_ARRAY;
         }
         var indices = new LinkedHashSet<String>();
-        forEachFromRelation(plan, relation -> indices.addAll(asList(Strings.commaDelimitedListToStringArray(relation.indexPattern()))));
+        forEachRelation(plan, relation -> indices.addAll(asList(Strings.commaDelimitedListToStringArray(relation.indexPattern()))));
         return indices.toArray(String[]::new);
     }
 
-    /**
-     * Iterates over the plan and applies the action to each {@link EsRelation} node.
-     * <p>
-     *     This method ignores the right side of joins.
-     * </p>
-     */
-    private static void forEachFromRelation(PhysicalPlan plan, Consumer<EsRelation> action) {
-        // Take the non-join-side fragments
-        forEachUpWithChildren(plan, FragmentExec.class, fragment -> {
-            // Take the non-join-side relations
-            forEachUpWithChildren(
-                fragment.fragment(),
-                EsRelation.class,
-                action,
-                node -> node instanceof Join join ? List.of(join.left()) : node.children()
-            );
-        }, node -> node instanceof LookupJoinExec join ? List.of(join.left()) : node.children());
-    }
-
-    /**
-     * Similar to {@link Node#forEachUp(Class, Consumer)}, but with a custom callback to get the node children.
-     */
-    private static <T extends Node<T>, E extends T> void forEachUpWithChildren(
-        T node,
-        Class<E> typeToken,
-        Consumer<? super E> action,
-        Function<? super T, Collection<T>> childrenGetter
-    ) {
-        childrenGetter.apply(node).forEach(c -> forEachUpWithChildren(c, typeToken, action, childrenGetter));
-        if (typeToken.isInstance(node)) {
-            action.accept(typeToken.cast(node));
-        }
+    private static void forEachRelation(PhysicalPlan plan, Consumer<EsRelation> action) {
+        plan.forEachDown(FragmentExec.class, f -> f.fragment().forEachDown(EsRelation.class, r -> {
+            if (r.indexMode() != IndexMode.LOOKUP) {
+                action.accept(r);
+            }
+        }));
     }
 
     public static PhysicalPlan localPlan(