diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java index 6d2df7046c..d10f0cde79 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleSet.java @@ -84,6 +84,7 @@ import org.apache.doris.nereids.rules.rewrite.PushdownAliasThroughJoin; import org.apache.doris.nereids.rules.rewrite.PushdownExpressionsInHashCondition; import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughAggregation; import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughJoin; +import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughPartitionTopN; import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughProject; import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughRepeat; import org.apache.doris.nereids.rules.rewrite.PushdownFilterThroughSetOperation; @@ -139,7 +140,8 @@ public class RuleSet { new MergeGenerates(), new MergeLimits(), new PushdownAliasThroughJoin(), - new PushdownFilterThroughWindow() + new PushdownFilterThroughWindow(), + new PushdownFilterThroughPartitionTopN() ); public static final List IMPLEMENTATION_RULES = planRuleFactories() 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 a1242c4a84..25d4db8abe 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 @@ -101,6 +101,7 @@ public enum RuleType { EXTRACT_AND_NORMALIZE_WINDOW_EXPRESSIONS(RuleTypeClass.REWRITE), CHECK_AND_STANDARDIZE_WINDOW_FUNCTION_AND_FRAME(RuleTypeClass.REWRITE), CHECK_MATCH_EXPRESSION(RuleTypeClass.REWRITE), + CREATE_PARTITION_TOPN_FOR_WINDOW(RuleTypeClass.REWRITE), AGGREGATE_DISASSEMBLE(RuleTypeClass.REWRITE), SIMPLIFY_AGG_GROUP_BY(RuleTypeClass.REWRITE), DISTINCT_AGGREGATE_DISASSEMBLE(RuleTypeClass.REWRITE), @@ -149,6 +150,7 @@ public enum RuleType { PUSHDOWN_FILTER_THROUGH_PROJECT(RuleTypeClass.REWRITE), PUSHDOWN_FILTER_THROUGH_PROJECT_UNDER_LIMIT(RuleTypeClass.REWRITE), PUSHDOWN_FILTER_THROUGH_WINDOW(RuleTypeClass.REWRITE), + PUSHDOWN_FILTER_THROUGH_PARTITION_TOPN(RuleTypeClass.REWRITE), PUSHDOWN_PROJECT_THROUGH_LIMIT(RuleTypeClass.REWRITE), PUSHDOWN_ALIAS_THROUGH_JOIN(RuleTypeClass.REWRITE), PUSHDOWN_ALIAS_INTO_UNION_ALL(RuleTypeClass.REWRITE), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CreatePartitionTopNFromWindow.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CreatePartitionTopNFromWindow.java index 1a4ae3ef1b..8a4d7a42d3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CreatePartitionTopNFromWindow.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CreatePartitionTopNFromWindow.java @@ -82,7 +82,10 @@ public class CreatePartitionTopNFromWindow extends OneRewriteRuleFactory { LogicalWindow window = filter.child(); // We have already done such optimization rule, so just ignore it. - if (window.child(0) instanceof LogicalPartitionTopN) { + if (window.child(0) instanceof LogicalPartitionTopN + || (window.child(0) instanceof LogicalFilter + && window.child(0).child(0) != null + && window.child(0).child(0) instanceof LogicalPartitionTopN)) { return filter; } @@ -137,7 +140,7 @@ public class CreatePartitionTopNFromWindow extends OneRewriteRuleFactory { return filter; } return filter.withChildren(newWindow.get()); - }).toRule(RuleType.PUSHDOWN_FILTER_THROUGH_WINDOW); + }).toRule(RuleType.CREATE_PARTITION_TOPN_FOR_WINDOW); } private Set extractRelatedConjuncts(Set conjuncts, ExprId slotRefID) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownFilterThroughPartitionTopN.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownFilterThroughPartitionTopN.java new file mode 100644 index 0000000000..da31fb2ae3 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushdownFilterThroughPartitionTopN.java @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +package org.apache.doris.nereids.rules.rewrite; + +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Expression; +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.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalPartitionTopN; + +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.ImmutableSet.Builder; + +import java.util.Set; + +/** + * Push down the 'filter' pass through the 'partitionTopN' if filter key is partitionTopN's partition key. + * Logical plan tree: + * any_node + * | + * filter (a <= 100) + * | + * partition topn (PARTITION BY a) + * | + * any_node + * transformed to: + * any_node + * | + * partition topn (PARTITION BY a) + * | + * filter (a <= 100) + * | + * any_node + */ + +public class PushdownFilterThroughPartitionTopN extends OneRewriteRuleFactory { + + @Override + public Rule build() { + return logicalFilter(logicalPartitionTopN()).thenApply(ctx -> { + LogicalFilter> filter = ctx.root; + LogicalPartitionTopN partitionTopN = filter.child(); + // follow the similar checking and transformation rule as + // PushdownFilterThroughWindow + Builder bottomConjunctsBuilder = ImmutableSet.builder(); + Builder upperConjunctsBuilder = ImmutableSet.builder(); + for (Expression expr : filter.getConjuncts()) { + boolean pushed = false; + Set exprInputSlots = expr.getInputSlots(); + for (Expression partitionKey : partitionTopN.getPartitionKeys()) { + if (partitionKey instanceof SlotReference + && exprInputSlots.size() == 1 + && partitionKey.getInputSlots().containsAll(exprInputSlots)) { + bottomConjunctsBuilder.add(expr); + pushed = true; + break; + } + } + if (!pushed) { + upperConjunctsBuilder.add(expr); + } + } + ImmutableSet bottomConjuncts = bottomConjunctsBuilder.build(); + ImmutableSet upperConjuncts = upperConjunctsBuilder.build(); + if (bottomConjuncts.isEmpty()) { + return null; + } + + LogicalFilter bottomFilter = new LogicalFilter<>(bottomConjuncts, partitionTopN.child()); + partitionTopN = (LogicalPartitionTopN) partitionTopN.withChildren(bottomFilter); + if (upperConjuncts.isEmpty()) { + return partitionTopN; + } else { + return filter.withConjunctsAndChild(upperConjuncts, partitionTopN); + } + }).toRule(RuleType.PUSHDOWN_FILTER_THROUGH_PARTITION_TOPN); + } + +} diff --git a/regression-test/data/nereids_syntax_p0/push_filter_through_ptopn.out b/regression-test/data/nereids_syntax_p0/push_filter_through_ptopn.out new file mode 100644 index 0000000000..05343ae859 --- /dev/null +++ b/regression-test/data/nereids_syntax_p0/push_filter_through_ptopn.out @@ -0,0 +1,52 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !1 -- +1 1 +1 2 + +-- !shape_1 -- +PhysicalResultSink +--PhysicalDistribute +----filter((rn <= 2)) +------PhysicalWindow +--------PhysicalQuickSort +----------PhysicalPartitionTopN +------------PhysicalProject +--------------filter((T.a = 1)) +----------------PhysicalOlapScan[push_filter_through_ptopn_tbl] + +-- !2 -- +1 1 +1 1 + +-- !shape_2 -- +PhysicalResultSink +--PhysicalDistribute +----PhysicalProject +------filter((rn <= 2)) +--------PhysicalWindow +----------PhysicalQuickSort +------------PhysicalPartitionTopN +--------------filter((T.a = 1)) +----------------PhysicalOlapScan[push_filter_through_ptopn_tbl] + +-- !3 -- +PhysicalResultSink +--PhysicalDistribute +----filter((T.b = 2) and (rn <= 2)) +------PhysicalWindow +--------PhysicalQuickSort +----------PhysicalPartitionTopN +------------PhysicalOlapScan[push_filter_through_ptopn_tbl] + +-- !4 -- +PhysicalResultSink +--PhysicalDistribute +----PhysicalProject +------filter((T.b = 2) and (rn <= 2)) +--------PhysicalWindow +----------PhysicalPartitionTopN +------------PhysicalDistribute +--------------PhysicalPartitionTopN +----------------PhysicalProject +------------------PhysicalOlapScan[push_filter_through_ptopn_tbl] + diff --git a/regression-test/suites/nereids_syntax_p0/push_filter_through_ptopn.groovy b/regression-test/suites/nereids_syntax_p0/push_filter_through_ptopn.groovy new file mode 100644 index 0000000000..162e71074f --- /dev/null +++ b/regression-test/suites/nereids_syntax_p0/push_filter_through_ptopn.groovy @@ -0,0 +1,53 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +suite("push_filter_through_ptopn") { + sql """set enable_nereids_planner=true""" + sql """set enable_partition_topn=true""" + sql """ + DROP TABLE IF EXISTS push_filter_through_ptopn_tbl + """ + sql """ + CREATE TABLE `push_filter_through_ptopn_tbl` ( + `a` int(11) NULL, + `b` int NULL + ) ENGINE=OLAP + duplicate KEY(`a`) + COMMENT 'OLAP' + DISTRIBUTED BY HASH(`a`) BUCKETS 4 + PROPERTIES ( + "replication_allocation" = "tag.location.default: 1", + "is_being_synced" = "false", + "storage_format" = "V2", + "light_schema_change" = "true", + "disable_auto_compaction" = "false", + "enable_single_replica_compaction" = "false" + ); + """ + sql """ + insert into push_filter_through_ptopn_tbl values (1, 2), (1, 3), (2, 4), (2, 5); + """ + + qt_1 """select * from (select a, row_number() over (partition by a) rn from push_filter_through_ptopn_tbl) T where a=1 and rn<=2;""" + qt_shape_1 """explain shape plan select * from (select a, row_number() over (partition by a) rn from push_filter_through_ptopn_tbl) T where a=1 and rn<=2;""" + qt_2 """select * from (select a, row_number() over (partition by b, a) rn from push_filter_through_ptopn_tbl) T where a=1 and rn <=2;""" + qt_shape_2 """explain shape plan select * from (select a, row_number() over (partition by b, a) rn from push_filter_through_ptopn_tbl) T where a=1 and rn<=2;""" + + qt_3 """explain shape plan select * from (select a, b, row_number() over (partition by a) rn from push_filter_through_ptopn_tbl) T where b=2 and rn<=2;""" + qt_4 """explain shape plan select * from (select a, b, row_number() over (partition by a+b) rn from push_filter_through_ptopn_tbl) T where b=2 and rn<=2;""" + +}