diff --git a/docs/changelog/123427.yaml b/docs/changelog/123427.yaml new file mode 100644 index 000000000000..50c29edb7972 --- /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 c5139d45f4b3..ff8833d4633a 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.aggregation.AggregatorMode; 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; @@ -21,7 +22,6 @@ 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,7 +33,6 @@ import org.elasticsearch.xpack.esql.optimizer.LocalPhysicalOptimizerContext; 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; @@ -41,7 +40,6 @@ 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.SearchContextStats; 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 class PlannerUtils { return Set.of(); } var indices = new LinkedHashSet(); - forEachFromRelation(plan, relation -> indices.addAll(relation.concreteIndices())); + forEachRelation(plan, relation -> indices.addAll(relation.concreteIndices())); return indices; } @@ -122,42 +118,16 @@ public class PlannerUtils { return Strings.EMPTY_ARRAY; } var indices = new LinkedHashSet(); - 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. - *

- * This method ignores the right side of joins. - *

- */ - private static void forEachFromRelation(PhysicalPlan plan, Consumer 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 , E extends T> void forEachUpWithChildren( - T node, - Class typeToken, - Consumer action, - Function> 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 action) { + plan.forEachDown(FragmentExec.class, f -> f.fragment().forEachDown(EsRelation.class, r -> { + if (r.indexMode() != IndexMode.LOOKUP) { + action.accept(r); + } + })); } public static PhysicalPlan localPlan(