From 0fb4b9716f85c34d9a82110439ddceeb28f76355 Mon Sep 17 00:00:00 2001
From: Costin Leau <costin@users.noreply.github.com>
Date: Wed, 26 Feb 2025 08:30:58 -0800
Subject: [PATCH] Update code for 8.17

---
 docs/changelog/123427.yaml                    |  5 +++
 .../xpack/esql/planner/PlannerUtils.java      | 32 ++++++++++++-------
 2 files changed, 26 insertions(+), 11 deletions(-)
 create mode 100644 docs/changelog/123427.yaml

diff --git a/docs/changelog/123427.yaml b/docs/changelog/123427.yaml
new file mode 100644
index 0000000000000..50c29edb79725
--- /dev/null
+++ b/docs/changelog/123427.yaml
@@ -0,0 +1,5 @@
+pr: 123427
+summary: Reduce iteration complexity for plan traversal
+area: ES|QL
+type: bug
+issues: []
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java
index c998af2215169..1a420144ee233 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/PlannerUtils.java
@@ -14,6 +14,7 @@
 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;
@@ -116,7 +117,7 @@ public static Set<String> planConcreteIndices(PhysicalPlan plan) {
             return Set.of();
         }
         var indices = new LinkedHashSet<String>();
-        plan.forEachUp(FragmentExec.class, f -> f.fragment().forEachUp(EsRelation.class, r -> indices.addAll(r.index().concreteIndices())));
+        forEachRelation(plan, relation -> indices.addAll(relation.concreteIndices()));
         return indices;
     }
 
@@ -128,21 +129,30 @@ public static String[] planOriginalIndices(PhysicalPlan plan) {
             return Strings.EMPTY_ARRAY;
         }
         var indices = new LinkedHashSet<String>();
-        plan.forEachUp(
-            FragmentExec.class,
-            f -> f.fragment()
-                .forEachUp(EsRelation.class, r -> indices.addAll(asList(Strings.commaDelimitedListToStringArray(r.index().name()))))
-        );
+        forEachRelation(plan, relation -> indices.addAll(asList(Strings.commaDelimitedListToStringArray(relation.indexPattern()))));
         return indices.toArray(String[]::new);
     }
 
-    public static PhysicalPlan localPlan(List<SearchExecutionContext> searchContexts, Configuration configuration, PhysicalPlan plan) {
-        return localPlan(configuration, plan, SearchContextStats.from(searchContexts));
+    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(
+        List<SearchExecutionContext> searchContexts,
+        Configuration configuration,
+        FoldContext foldCtx,
+        PhysicalPlan plan
+    ) {
+        return localPlan(configuration, foldCtx, plan, SearchContextStats.from(searchContexts));
     }
 
-    public static PhysicalPlan localPlan(Configuration configuration, PhysicalPlan plan, SearchStats searchStats) {
-        final var logicalOptimizer = new LocalLogicalPlanOptimizer(new LocalLogicalOptimizerContext(configuration, searchStats));
-        var physicalOptimizer = new LocalPhysicalPlanOptimizer(new LocalPhysicalOptimizerContext(configuration, searchStats));
+    public static PhysicalPlan localPlan(Configuration configuration, FoldContext foldCtx, PhysicalPlan plan, SearchStats searchStats) {
+        final var logicalOptimizer = new LocalLogicalPlanOptimizer(new LocalLogicalOptimizerContext(configuration, foldCtx, searchStats));
+        var physicalOptimizer = new LocalPhysicalPlanOptimizer(new LocalPhysicalOptimizerContext(configuration, foldCtx, searchStats));
 
         return localPlan(plan, logicalOptimizer, physicalOptimizer);
     }