Merge pull request #5924 from leiysky/constant-folding

feature(optimizer): Support constant folding

query/src/sql/optimizer/heuristic/mod.rs

@@ -26,6 +26,7 @@ use crate::sql::optimizer::SExpr;
 
 lazy_static! {
     pub static ref DEFAULT_REWRITE_RULES: Vec<RuleID> = vec![
+        RuleID::NormalizeScalarFilter,
         RuleID::EliminateFilter,
         RuleID::EliminateEvalScalar,
         RuleID::EliminateProject,

query/src/sql/optimizer/rule/factory.rs

@@ -15,6 +15,7 @@
 use common_exception::Result;
 
 use super::rewrite::RuleEliminateEvalScalar;
+use super::rewrite::RuleNormalizeScalarFilter;
 use super::rewrite::RulePushDownFilterEvalScalar;
 use super::rewrite::RulePushDownFilterJoin;
 use super::rewrite::RulePushDownFilterProject;
@@ -48,6 +49,7 @@ impl RuleFactory {
             RuleID::MergeProject => Ok(Box::new(RuleMergeProject::new())),
             RuleID::MergeEvalScalar => Ok(Box::new(RuleMergeEvalScalar::new())),
             RuleID::MergeFilter => Ok(Box::new(RuleMergeFilter::new())),
+            RuleID::NormalizeScalarFilter => Ok(Box::new(RuleNormalizeScalarFilter::new())),
         }
     }
 }

query/src/sql/optimizer/rule/mod.rs

@@ -43,6 +43,8 @@ pub trait Rule {
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
 pub enum RuleID {
+    // Rewrite rules
+    NormalizeScalarFilter,
     PushDownFilterProject,
     PushDownFilterEvalScalar,
     PushDownFilterJoin,
@@ -53,6 +55,7 @@ pub enum RuleID {
     MergeEvalScalar,
     MergeFilter,
 
+    // Implementation rules
     ImplementGet,
     ImplementHashJoin,
 }
@@ -71,6 +74,7 @@ impl Display for RuleID {
             RuleID::MergeProject => write!(f, "MergeProject"),
             RuleID::MergeEvalScalar => write!(f, "MergeEvalScalar"),
             RuleID::MergeFilter => write!(f, "MergeFilter"),
+            RuleID::NormalizeScalarFilter => write!(f, "NormalizeScalarFilter"),
         }
     }
 }

query/src/sql/optimizer/rule/rewrite/mod.rs

@@ -18,6 +18,7 @@ mod rule_eliminate_project;
 mod rule_merge_eval_scalar;
 mod rule_merge_filter;
 mod rule_merge_project;
+mod rule_normalize_scalar;
 mod rule_push_down_filter_eval_scalar;
 mod rule_push_down_filter_join;
 mod rule_push_down_filter_project;
@@ -28,6 +29,7 @@ pub use rule_eliminate_project::RuleEliminateProject;
 pub use rule_merge_eval_scalar::RuleMergeEvalScalar;
 pub use rule_merge_filter::RuleMergeFilter;
 pub use rule_merge_project::RuleMergeProject;
+pub use rule_normalize_scalar::RuleNormalizeScalarFilter;
 pub use rule_push_down_filter_eval_scalar::RulePushDownFilterEvalScalar;
 pub use rule_push_down_filter_join::RulePushDownFilterJoin;
 pub use rule_push_down_filter_project::RulePushDownFilterProject;

query/src/sql/optimizer/rule/rewrite/rule_normalize_scalar.rs

@@ -0,0 +1,130 @@
+// 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 common_datavalues::BooleanType;
+use common_datavalues::DataValue;
+use common_exception::Result;
+
+use crate::sql::optimizer::rule::Rule;
+use crate::sql::optimizer::RuleID;
+use crate::sql::optimizer::SExpr;
+use crate::sql::plans::ConstantExpr;
+use crate::sql::plans::Filter;
+use crate::sql::plans::PatternPlan;
+use crate::sql::plans::RelOp;
+use crate::sql::plans::Scalar;
+
+fn normalize_predicates(predicates: Vec<Scalar>) -> Vec<Scalar> {
+    [remove_true_predicate, normalize_falsy_predicate]
+        .into_iter()
+        .fold(predicates, |acc, f| f(acc))
+}
+
+fn is_true(predicate: &Scalar) -> bool {
+    matches!(
+        predicate,
+        Scalar::ConstantExpr(ConstantExpr {
+            value: DataValue::Boolean(true),
+            ..
+        })
+    )
+}
+
+fn is_falsy(predicate: &Scalar) -> bool {
+    matches!(
+        predicate,
+        Scalar::ConstantExpr(ConstantExpr {
+            value,
+            ..
+        }) if value == &DataValue::Boolean(false) || value == &DataValue::Null
+    )
+}
+
+fn remove_true_predicate(predicates: Vec<Scalar>) -> Vec<Scalar> {
+    predicates.into_iter().filter(|p| !is_true(p)).collect()
+}
+
+fn normalize_falsy_predicate(predicates: Vec<Scalar>) -> Vec<Scalar> {
+    if predicates.iter().any(is_falsy) {
+        vec![ConstantExpr {
+            value: DataValue::Boolean(false),
+            data_type: BooleanType::new_impl(),
+        }
+        .into()]
+    } else {
+        predicates
+    }
+}
+
+/// Rule to normalize a Filter, including:
+/// - Remove true predicates
+/// - If there is a NULL or FALSE conjuction, replace the
+/// whole filter with FALSE
+pub struct RuleNormalizeScalarFilter {
+    id: RuleID,
+    pattern: SExpr,
+}
+
+impl RuleNormalizeScalarFilter {
+    pub fn new() -> Self {
+        Self {
+            id: RuleID::NormalizeScalarFilter,
+            // Filter
+            //  \
+            //   *
+            pattern: SExpr::create_unary(
+                PatternPlan {
+                    plan_type: RelOp::Filter,
+                }
+                .into(),
+                SExpr::create_leaf(
+                    PatternPlan {
+                        plan_type: RelOp::Pattern,
+                    }
+                    .into(),
+                ),
+            ),
+        }
+    }
+}
+
+impl Rule for RuleNormalizeScalarFilter {
+    fn id(&self) -> RuleID {
+        self.id
+    }
+
+    fn apply(
+        &self,
+        s_expr: &SExpr,
+        state: &mut crate::sql::optimizer::rule::TransformState,
+    ) -> Result<()> {
+        let mut filter: Filter = s_expr.plan().clone().try_into()?;
+
+        if filter
+            .predicates
+            .iter()
+            .any(|p| is_true(p) || (is_falsy(p) && filter.predicates.len() > 1))
+        {
+            filter.predicates = normalize_predicates(filter.predicates);
+            state.add_result(SExpr::create_unary(filter.into(), s_expr.child(0)?.clone()));
+            Ok(())
+        } else {
+            Ok(())
+        }
+    }
+
+    fn pattern(&self) -> &SExpr {
+        &self.pattern
+    }
+}