This is an automated email from the ASF dual-hosted git repository. huajianlan pushed a commit to branch nested_column_prune in repository https://gitbox.apache.org/repos/asf/doris.git
commit e54ee541bb31a8babdb5fb1e234a9355a91ea0dd Author: 924060929 <[email protected]> AuthorDate: Wed Oct 22 22:33:29 2025 +0800 push down project through partition top n --- .../org/apache/doris/nereids/rules/RuleType.java | 2 ++ .../rules/rewrite/PushDownNestedColumn.java | 42 ++++++++-------------- .../rules/rewrite/PruneNestedColumnTest.java | 40 ++++++++++++++++++--- 3 files changed, 52 insertions(+), 32 deletions(-) diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index 0ca77df2f9c..fec9e437579 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -198,6 +198,8 @@ public enum RuleType { PUSH_DOWN_NESTED_COLUMN_THROUGH_JOIN(RuleTypeClass.REWRITE), PUSH_DOWN_NESTED_COLUMN_THROUGH_UNION(RuleTypeClass.REWRITE), PUSH_DOWN_NESTED_COLUMN_THROUGH_WINDOW(RuleTypeClass.REWRITE), + PUSH_DOWN_NESTED_COLUMN_THROUGH_PARTITION_TOP_N(RuleTypeClass.REWRITE), + PUSH_DOWN_NESTED_COLUMN_THROUGH_DEFER_MATERIALIZE_TOP_N(RuleTypeClass.REWRITE), // Pushdown filter PUSH_DOWN_FILTER_THROUGH_JOIN(RuleTypeClass.REWRITE), PUSH_DOWN_FILTER_THROUGH_LEFT_SEMI_JOIN(RuleTypeClass.REWRITE), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownNestedColumn.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownNestedColumn.java index d728832d74c..ff4da86d783 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownNestedColumn.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownNestedColumn.java @@ -30,11 +30,10 @@ import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.algebra.SetOperation.Qualifier; -import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; import org.apache.doris.nereids.trees.plans.logical.LogicalProject; import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; -import org.apache.doris.nereids.trees.plans.logical.LogicalWindow; import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.ImmutableCollection; @@ -57,11 +56,17 @@ public class PushDownNestedColumn implements RewriteRuleFactory { public List<Rule> buildRules() { return ImmutableList.of( RuleType.PUSH_DOWN_NESTED_COLUMN_THROUGH_JOIN.build( - logicalProject(logicalJoin()).thenApply(this::pushThroughJoin) + logicalProject(logicalJoin()).thenApply(this::defaultPushDown) ), RuleType.PUSH_DOWN_NESTED_COLUMN_THROUGH_WINDOW.build( - logicalProject(logicalWindow()).thenApply(this::pushThroughWindow) + logicalProject(logicalWindow()).thenApply(this::defaultPushDown) ), + RuleType.PUSH_DOWN_NESTED_COLUMN_THROUGH_PARTITION_TOP_N.build( + logicalProject(logicalPartitionTopN()).thenApply(this::defaultPushDown) + ), + // RuleType.PUSH_DOWN_NESTED_COLUMN_THROUGH_DEFER_MATERIALIZE_TOP_N.build( + // logicalProject(logicalDeferMaterializeTopN()).thenApply(this::defaultPushDown) + // ), RuleType.PUSH_DOWN_NESTED_COLUMN_THROUGH_UNION.build( logicalProject( logicalUnion().when(u -> u.getQualifier() == Qualifier.ALL) @@ -70,30 +75,11 @@ public class PushDownNestedColumn implements RewriteRuleFactory { ); } - private Plan pushThroughWindow(MatchingContext<LogicalProject<LogicalWindow<Plan>>> ctx) { - LogicalProject<LogicalWindow<Plan>> project = ctx.root; - LogicalWindow<Plan> window = project.child(); - PushdownProjectHelper pushdownProjectHelper - = new PushdownProjectHelper(ctx.statementContext, window); - - Pair<Boolean, List<NamedExpression>> pushProjects - = pushdownProjectHelper.pushDownExpressions(project.getProjects()); - - if (pushProjects.first) { - List<Plan> newWindowChildren = pushdownProjectHelper.buildNewChildren(); - return new LogicalProject<>( - pushProjects.second, - window.withChildren(newWindowChildren) - ); - } - return project; - } - - private Plan pushThroughJoin(MatchingContext<LogicalProject<LogicalJoin<Plan, Plan>>> ctx) { - LogicalProject<LogicalJoin<Plan, Plan>> project = ctx.root; - LogicalJoin<Plan, Plan> join = project.child(); + private <C extends LogicalPlan> Plan defaultPushDown(MatchingContext<LogicalProject<C>> ctx) { + LogicalProject<C> project = ctx.root; + C child = project.child(); PushdownProjectHelper pushdownProjectHelper - = new PushdownProjectHelper(ctx.statementContext, join); + = new PushdownProjectHelper(ctx.statementContext, child); Pair<Boolean, List<NamedExpression>> pushProjects = pushdownProjectHelper.pushDownExpressions(project.getProjects()); @@ -102,7 +88,7 @@ public class PushDownNestedColumn implements RewriteRuleFactory { List<Plan> newJoinChildren = pushdownProjectHelper.buildNewChildren(); return new LogicalProject<>( pushProjects.second, - join.withChildren(newJoinChildren) + child.withChildren(newJoinChildren) ); } return project; diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java index 080eb13eed1..ec8f26e1a43 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/PruneNestedColumnTest.java @@ -274,8 +274,8 @@ public class PruneNestedColumnTest extends TestWithFeService implements MemoPatt @Test public void testUnion() throws Throwable { assertColumn("select struct_element(s, 'city') from (select s from tbl union all select null)a", - "struct<city:text,data:array<map<int,struct<a:int,b:double>>>>", - ImmutableList.of(path("s")), + "struct<city:text>", + ImmutableList.of(path("s", "city")), ImmutableList.of() ); @@ -289,8 +289,8 @@ public class PruneNestedColumnTest extends TestWithFeService implements MemoPatt @Test public void testCteAndUnion() throws Throwable { assertColumn("with t as (select id, s from tbl) select struct_element(s, 'city') from (select * from t union all select 1, null) tmp", - "struct<city:text,data:array<map<int,struct<a:int,b:double>>>>", - ImmutableList.of(path("s")), + "struct<city:text>", + ImmutableList.of(path("s", "city")), ImmutableList.of() ); @@ -356,6 +356,38 @@ public class PruneNestedColumnTest extends TestWithFeService implements MemoPatt ); } + @Test + public void testPushDownThroughPartitionTopN() { + PlanChecker.from(connectContext) + .analyze("select struct_element(s, 'city'), r from (select s, rank() over(partition by id) r from tbl t limit 10)a") + .rewrite() + .matches( + logicalResultSink( + logicalLimit( + logicalLimit( + logicalProject( + logicalWindow( + logicalPartitionTopN( + logicalProject( + logicalOlapScan() + ).when(p -> { + Assertions.assertEquals(2, p.getProjects().size()); + Assertions.assertTrue(p.getProjects().stream() + .anyMatch(o -> o instanceof Alias && o.child(0) instanceof StructElement)); + return true; + }) + ) + ) + ).when(p -> { + Assertions.assertTrue(p.getProjects().size() == 2 && p.getProjects().get(0) instanceof SlotReference); + return true; + }) + ) + ) + ) + ); + } + @Test public void testPushDownThroughUnion() { PlanChecker.from(connectContext) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
