Merge pull request #9209 from xudong963/filter_push_down_union

BohuTANG · web-flow · commit 288c3bba04b7 · 2022-12-12T20:37:46.000+08:00
feat: push down `filter` to `union`
diff --git a/src/query/sql/src/planner/optimizer/heuristic/heuristic.rs b/src/query/sql/src/planner/optimizer/heuristic/heuristic.rs
@@ -38,6 +38,7 @@ pub static DEFAULT_REWRITE_RULES: Lazy<Vec<RuleID>> = Lazy::new(|| {
         RuleID::EliminateEvalScalar,
         RuleID::MergeFilter,
         RuleID::MergeEvalScalar,
+        RuleID::PushDownFilterUnion,
         RuleID::PushDownLimitSort,
         RuleID::PushDownLimitOuterJoin,
         RuleID::PushDownLimitScan,
diff --git a/src/query/sql/src/planner/optimizer/rule/factory.rs b/src/query/sql/src/planner/optimizer/rule/factory.rs
@@ -27,6 +27,7 @@ use crate::optimizer::rule::rewrite::RuleEliminateFilter;
 use crate::optimizer::rule::rewrite::RuleMergeEvalScalar;
 use crate::optimizer::rule::rewrite::RuleMergeFilter;
 use crate::optimizer::rule::rewrite::RulePushDownFilterScan;
+use crate::optimizer::rule::rewrite::RulePushDownFilterUnion;
 use crate::optimizer::rule::rewrite::RulePushDownLimitOuterJoin;
 use crate::optimizer::rule::rewrite::RulePushDownLimitScan;
 use crate::optimizer::rule::rewrite::RulePushDownLimitSort;
@@ -49,6 +50,7 @@ impl RuleFactory {
             RuleID::ImplementGet => Ok(Box::new(RuleImplementGet::new())),
             RuleID::ImplementHashJoin => Ok(Box::new(RuleImplementHashJoin::new())),
             RuleID::EliminateEvalScalar => Ok(Box::new(RuleEliminateEvalScalar::new())),
+            RuleID::PushDownFilterUnion => Ok(Box::new(RulePushDownFilterUnion::new())),
             RuleID::PushDownFilterEvalScalar => Ok(Box::new(RulePushDownFilterEvalScalar::new())),
             RuleID::PushDownFilterJoin => Ok(Box::new(RulePushDownFilterJoin::new())),
             RuleID::PushDownFilterScan => Ok(Box::new(RulePushDownFilterScan::new())),
diff --git a/src/query/sql/src/planner/optimizer/rule/rewrite/mod.rs b/src/query/sql/src/planner/optimizer/rule/rewrite/mod.rs
@@ -22,6 +22,7 @@ mod rule_normalize_scalar;
 mod rule_push_down_filter_eval_scalar;
 mod rule_push_down_filter_join;
 mod rule_push_down_filter_scan;
+mod rule_push_down_filter_union;
 mod rule_push_down_limit_join;
 mod rule_push_down_limit_scan;
 mod rule_push_down_limit_sort;
@@ -39,6 +40,7 @@ pub use rule_push_down_filter_eval_scalar::RulePushDownFilterEvalScalar;
 pub use rule_push_down_filter_join::try_push_down_filter_join;
 pub use rule_push_down_filter_join::RulePushDownFilterJoin;
 pub use rule_push_down_filter_scan::RulePushDownFilterScan;
+pub use rule_push_down_filter_union::RulePushDownFilterUnion;
 pub use rule_push_down_limit_join::RulePushDownLimitOuterJoin;
 pub use rule_push_down_limit_scan::RulePushDownLimitScan;
 pub use rule_push_down_limit_sort::RulePushDownLimitSort;
diff --git a/src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_filter_union.rs b/src/query/sql/src/planner/optimizer/rule/rewrite/rule_push_down_filter_union.rs
@@ -0,0 +1,198 @@
+// Copyright 2022 Datafuse Labs.
+//
+// Licensed 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.
+
+use ahash::HashMap;
+use common_exception::ErrorCode;
+use common_exception::Result;
+
+use crate::optimizer::rule::Rule;
+use crate::optimizer::rule::TransformResult;
+use crate::optimizer::RuleID;
+use crate::optimizer::SExpr;
+use crate::plans::AggregateFunction;
+use crate::plans::AndExpr;
+use crate::plans::BoundColumnRef;
+use crate::plans::CastExpr;
+use crate::plans::ComparisonExpr;
+use crate::plans::Filter;
+use crate::plans::FunctionCall;
+use crate::plans::OrExpr;
+use crate::plans::PatternPlan;
+use crate::plans::RelOp;
+use crate::plans::Scalar;
+use crate::plans::UnionAll;
+use crate::ColumnBinding;
+use crate::IndexType;
+use crate::Visibility;
+
+// For a union query, it's not allowed to add `filter` after union
+// Such as: `(select * from t1 union all select * from t2) where a > 1`, it's invalid.
+// However, it's possible to have `filter` after `union` when involved `view`
+// Such as: `create view v_t as (select * from t1 union all select * from t2)`.
+// Then use the view with filter, `select * from v_t where a > 1`;
+// So it'll be efficient to push down `filter` to `union`, reduce the size of data to pull from table.
+pub struct RulePushDownFilterUnion {
+    id: RuleID,
+    pattern: SExpr,
+}
+
+impl RulePushDownFilterUnion {
+    pub fn new() -> Self {
+        Self {
+            id: RuleID::PushDownFilterUnion,
+            // Filter
+            //  \
+            //   UnionAll
+            //     /  \
+            //   ...   ...
+            pattern: SExpr::create_unary(
+                PatternPlan {
+                    plan_type: RelOp::Filter,
+                }
+                .into(),
+                SExpr::create_binary(
+                    PatternPlan {
+                        plan_type: RelOp::UnionAll,
+                    }
+                    .into(),
+                    SExpr::create_leaf(
+                        PatternPlan {
+                            plan_type: RelOp::Pattern,
+                        }
+                        .into(),
+                    ),
+                    SExpr::create_leaf(
+                        PatternPlan {
+                            plan_type: RelOp::Pattern,
+                        }
+                        .into(),
+                    ),
+                ),
+            ),
+        }
+    }
+}
+
+impl Rule for RulePushDownFilterUnion {
+    fn id(&self) -> RuleID {
+        self.id
+    }
+
+    fn apply(&self, s_expr: &SExpr, state: &mut TransformResult) -> Result<()> {
+        let filter: Filter = s_expr.plan().clone().try_into()?;
+        let union_s_expr = s_expr.child(0)?;
+        let union: UnionAll = union_s_expr.plan().clone().try_into()?;
+
+        // Create a filter which matches union's right child.
+        let index_pairs: HashMap<IndexType, IndexType> =
+            union.pairs.iter().map(|pair| (pair.0, pair.1)).collect();
+        let new_predicates = filter
+            .predicates
+            .iter()
+            .map(|predicate| replace_column_binding(&index_pairs, predicate.clone()))
+            .collect::<Result<Vec<_>>>()?;
+        let right_filer = Filter {
+            predicates: new_predicates,
+            is_having: filter.is_having,
+        };
+
+        let mut union_left_child = union_s_expr.child(0)?.clone();
+        let mut union_right_child = union_s_expr.child(1)?.clone();
+
+        // Add filter to union children
+        union_left_child = SExpr::create_unary(filter.into(), union_left_child);
+        union_right_child = SExpr::create_unary(right_filer.into(), union_right_child);
+
+        let result = SExpr::create_binary(union.into(), union_left_child, union_right_child);
+        state.add_result(result);
+
+        Ok(())
+    }
+
+    fn pattern(&self) -> &SExpr {
+        &self.pattern
+    }
+}
+
+fn replace_column_binding(
+    index_pairs: &HashMap<IndexType, IndexType>,
+    scalar: Scalar,
+) -> Result<Scalar> {
+    match scalar {
+        Scalar::BoundColumnRef(column) => {
+            let index = column.column.index;
+            if index_pairs.contains_key(&index) {
+                let new_column = ColumnBinding {
+                    database_name: None,
+                    table_name: None,
+                    column_name: "".to_string(),
+                    index: *index_pairs.get(&index).unwrap(),
+                    data_type: column.column.data_type,
+                    visibility: Visibility::Visible,
+                };
+                return Ok(Scalar::BoundColumnRef(BoundColumnRef {
+                    column: new_column,
+                }));
+            }
+            Ok(Scalar::BoundColumnRef(column))
+        }
+        constant_expr @ Scalar::ConstantExpr(_) => Ok(constant_expr),
+        Scalar::AndExpr(expr) => Ok(Scalar::AndExpr(AndExpr {
+            left: Box::new(replace_column_binding(index_pairs, *expr.left)?),
+            right: Box::new(replace_column_binding(index_pairs, *expr.right)?),
+            return_type: expr.return_type,
+        })),
+        Scalar::OrExpr(expr) => Ok(Scalar::OrExpr(OrExpr {
+            left: Box::new(replace_column_binding(index_pairs, *expr.left)?),
+            right: Box::new(replace_column_binding(index_pairs, *expr.right)?),
+            return_type: expr.return_type,
+        })),
+        Scalar::ComparisonExpr(expr) => Ok(Scalar::ComparisonExpr(ComparisonExpr {
+            op: expr.op,
+            left: Box::new(replace_column_binding(index_pairs, *expr.left)?),
+            right: Box::new(replace_column_binding(index_pairs, *expr.right)?),
+            return_type: expr.return_type,
+        })),
+        Scalar::AggregateFunction(expr) => Ok(Scalar::AggregateFunction(AggregateFunction {
+            display_name: expr.display_name,
+            func_name: expr.func_name,
+            distinct: expr.distinct,
+            params: expr.params,
+            args: expr
+                .args
+                .into_iter()
+                .map(|arg| replace_column_binding(index_pairs, arg))
+                .collect::<Result<Vec<_>>>()?,
+            return_type: expr.return_type,
+        })),
+        Scalar::FunctionCall(expr) => Ok(Scalar::FunctionCall(FunctionCall {
+            arguments: expr
+                .arguments
+                .into_iter()
+                .map(|arg| replace_column_binding(index_pairs, arg))
+                .collect::<Result<Vec<_>>>()?,
+            func_name: expr.func_name,
+            arg_types: expr.arg_types,
+            return_type: expr.return_type,
+        })),
+        Scalar::CastExpr(expr) => Ok(Scalar::CastExpr(CastExpr {
+            argument: Box::new(replace_column_binding(index_pairs, *(expr.argument))?),
+            from_type: expr.from_type,
+            target_type: expr.target_type,
+        })),
+        Scalar::SubqueryExpr(_) => Err(ErrorCode::Unimplemented(
+            "replace_column_binding: don't support subquery",
+        )),
+    }
+}
diff --git a/src/query/sql/src/planner/optimizer/rule/rule.rs b/src/query/sql/src/planner/optimizer/rule/rule.rs
@@ -36,6 +36,7 @@ pub enum RuleID {
     NormalizeScalarFilter,
     NormalizeDisjunctiveFilter,
     PushDownFilterEvalScalar,
+    PushDownFilterUnion,
     PushDownFilterJoin,
     PushDownFilterScan,
     PushDownLimitOuterJoin,
@@ -62,6 +63,7 @@ pub enum RuleID {
 impl Display for RuleID {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
+            RuleID::PushDownFilterUnion => write!(f, "PushDownFilterUnion"),
             RuleID::PushDownFilterEvalScalar => write!(f, "PushDownFilterEvalScalar"),
             RuleID::PushDownFilterJoin => write!(f, "PushDownFilterJoin"),
             RuleID::PushDownFilterScan => write!(f, "PushDownFilterScan"),
diff --git a/tests/logictest/suites/query/union.test b/tests/logictest/suites/query/union.test
@@ -144,4 +144,77 @@ ORDER BY age,
 17-23  2
 24-30  2
 年轻  2
-非年轻  2
+非年轻  2
+
+statement ok
+drop view if exists t;
+
+statement ok
+drop table if exists t1;
+
+statement ok
+drop table if exists t2;
+
+statement ok
+create table t1 (a int, b int);
+
+statement ok
+insert into t1 values (1, 2), (2, 3);
+
+statement ok
+create table t2 (a int, b int);
+
+statement ok
+insert into t2 values (1, 2), (2, 3);
+
+statement ok
+create view t as select * from t1 union all select * from t2;
+
+statement query T
+explain select * from t where a > b;
+
+----
+UnionAll
+├── TableScan
+│   ├── table: default.default.t1
+│   ├── read rows: 0
+│   ├── read bytes: 0
+│   ├── partitions total: 1
+│   ├── partitions scanned: 0
+│   └── push downs: [filters: [(a > b)], limit: NONE]
+└── TableScan
+    ├── table: default.default.t2
+    ├── read rows: 0
+    ├── read bytes: 0
+    ├── partitions total: 1
+    ├── partitions scanned: 0
+    └── push downs: [filters: [(a > b)], limit: NONE]
+
+statement query T
+explain select * from t where a > 1;
+
+----
+UnionAll
+├── TableScan
+│   ├── table: default.default.t1
+│   ├── read rows: 2
+│   ├── read bytes: 62
+│   ├── partitions total: 1
+│   ├── partitions scanned: 1
+│   └── push downs: [filters: [(a > 1)], limit: NONE]
+└── TableScan
+    ├── table: default.default.t2
+    ├── read rows: 2
+    ├── read bytes: 62
+    ├── partitions total: 1
+    ├── partitions scanned: 1
+    └── push downs: [filters: [(a > 1)], limit: NONE]
+
+statement ok
+drop table t1;
+
+statement ok
+drop table t2;
+
+statement ok
+drop view t;
