You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/01/26 11:36:12 UTC
[GitHub] [arrow] alamb commented on a change in pull request #9309: ARROW-11366: [Datafusion] support boolean literal in comparison expression
alamb commented on a change in pull request #9309:
URL: https://github.com/apache/arrow/pull/9309#discussion_r564440788
##########
File path: rust/datafusion/src/optimizer/boolean_comparison.rs
##########
@@ -0,0 +1,270 @@
+// Licensed to the Apache Software Foundation (ASF) under one
Review comment:
I suggest calling this module 'constant_folding.rs` and the optimization `ConstantFolding` -- while this particular PR has code to fold boolean constants, the idea is much more general (e.g you might imagine folding things like `A < (5+3)` into `A < 8` and that sort of thing).
##########
File path: rust/datafusion/src/optimizer/boolean_comparison.rs
##########
@@ -0,0 +1,270 @@
+// 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.
+
+//! Boolean comparision rule rewrites redudant comparison expression involing boolean literal into
+//! unary expression.
+
+use std::sync::Arc;
+
+use crate::error::Result;
+use crate::logical_plan::{Expr, LogicalPlan, Operator};
+use crate::optimizer::optimizer::OptimizerRule;
+use crate::optimizer::utils;
+use crate::scalar::ScalarValue;
+
+/// Optimizer that simplifies comparison expressions involving boolean literals.
+///
+/// Recursively go through all expressionss and simplify the following cases:
+/// * `expr = ture` to `expr`
+/// * `expr = false` to `!expr`
+pub struct BooleanComparison {}
+
+impl BooleanComparison {
+ #[allow(missing_docs)]
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl OptimizerRule for BooleanComparison {
+ fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
+ match plan {
+ LogicalPlan::Filter { predicate, input } => Ok(LogicalPlan::Filter {
+ predicate: optimize_expr(predicate),
+ input: Arc::new(self.optimize(input)?),
+ }),
+ // Rest: recurse into plan, apply optimization where possible
+ LogicalPlan::Projection { .. }
+ | LogicalPlan::Aggregate { .. }
+ | LogicalPlan::Limit { .. }
+ | LogicalPlan::Repartition { .. }
+ | LogicalPlan::CreateExternalTable { .. }
+ | LogicalPlan::Extension { .. }
+ | LogicalPlan::Sort { .. }
+ | LogicalPlan::Explain { .. }
+ | LogicalPlan::Join { .. } => {
+ let expr = utils::expressions(plan);
+
+ // apply the optimization to all inputs of the plan
+ let inputs = utils::inputs(plan);
+ let new_inputs = inputs
+ .iter()
+ .map(|plan| self.optimize(plan))
+ .collect::<Result<Vec<_>>>()?;
+
+ utils::from_plan(plan, &expr, &new_inputs)
+ }
+ LogicalPlan::TableScan { .. } | LogicalPlan::EmptyRelation { .. } => {
+ Ok(plan.clone())
+ }
+ }
+ }
+
+ fn name(&self) -> &str {
+ "boolean_comparison"
+ }
+}
+
+/// Recursively transverses the logical plan.
+fn optimize_expr(e: &Expr) -> Expr {
+ match e {
+ Expr::BinaryExpr { left, op, right } => {
+ let left = optimize_expr(left);
+ let right = optimize_expr(right);
+ match op {
+ Operator::Eq => match (&left, &right) {
+ (Expr::Literal(ScalarValue::Boolean(b)), _) => match b {
+ Some(true) => right,
+ Some(false) | None => Expr::Not(Box::new(right)),
+ },
+ (_, Expr::Literal(ScalarValue::Boolean(b))) => match b {
+ Some(true) => left,
+ Some(false) | None => Expr::Not(Box::new(left)),
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: Operator::Eq,
+ right: Box::new(right),
+ },
+ },
+ Operator::NotEq => match (&left, &right) {
+ (Expr::Literal(ScalarValue::Boolean(b)), _) => match b {
+ Some(false) | None => right,
+ Some(true) => Expr::Not(Box::new(right)),
+ },
+ (_, Expr::Literal(ScalarValue::Boolean(b))) => match b {
+ Some(false) | None => left,
+ Some(true) => Expr::Not(Box::new(left)),
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: Operator::NotEq,
+ right: Box::new(right),
+ },
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: op.clone(),
+ right: Box::new(right),
+ },
+ }
+ }
+ Expr::Not(expr) => Expr::Not(Box::new(optimize_expr(&expr))),
+ Expr::Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ } => {
+ if expr.is_none() {
+ // recurse into CASE WHEN condition expressions
+ Expr::Case {
+ expr: None,
+ when_then_expr: when_then_expr
+ .iter()
+ .map(|(when, then)| (Box::new(optimize_expr(when)), then.clone()))
+ .collect(),
+ else_expr: else_expr.clone(),
+ }
+ } else {
+ // when base expression is specified, when_then_expr conditions are literal values
+ // so we can just skip this case
+ e.clone()
+ }
+ }
+ Expr::Alias { .. }
+ | Expr::Negative { .. }
+ | Expr::Column { .. }
+ | Expr::InList { .. }
+ | Expr::IsNotNull { .. }
+ | Expr::IsNull { .. }
+ | Expr::Cast { .. }
+ | Expr::ScalarVariable { .. }
+ | Expr::Between { .. }
+ | Expr::Literal { .. }
+ | Expr::ScalarFunction { .. }
+ | Expr::ScalarUDF { .. }
+ | Expr::AggregateFunction { .. }
+ | Expr::AggregateUDF { .. }
+ | Expr::Sort { .. }
+ | Expr::Wildcard => e.clone(),
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::logical_plan::{col, lit, LogicalPlanBuilder};
+ use crate::test::*;
+
+ fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
+ let mut rule = BooleanComparison::new();
+ let optimized_plan = rule.optimize(plan).expect("failed to optimize plan");
+ let formatted_plan = format!("{:?}", optimized_plan);
+ assert_eq!(formatted_plan, expected);
+ }
+
+ #[test]
+ fn simplify_eq_expr() -> Result<()> {
Review comment:
I think adding a test for rewriting expressions in a non-filter plan would be valuable (e.g. make a join plan or something)
##########
File path: rust/datafusion/src/optimizer/boolean_comparison.rs
##########
@@ -0,0 +1,270 @@
+// 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.
+
+//! Boolean comparision rule rewrites redudant comparison expression involing boolean literal into
+//! unary expression.
+
+use std::sync::Arc;
+
+use crate::error::Result;
+use crate::logical_plan::{Expr, LogicalPlan, Operator};
+use crate::optimizer::optimizer::OptimizerRule;
+use crate::optimizer::utils;
+use crate::scalar::ScalarValue;
+
+/// Optimizer that simplifies comparison expressions involving boolean literals.
+///
+/// Recursively go through all expressionss and simplify the following cases:
+/// * `expr = ture` to `expr`
+/// * `expr = false` to `!expr`
+pub struct BooleanComparison {}
+
+impl BooleanComparison {
+ #[allow(missing_docs)]
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl OptimizerRule for BooleanComparison {
+ fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
+ match plan {
+ LogicalPlan::Filter { predicate, input } => Ok(LogicalPlan::Filter {
+ predicate: optimize_expr(predicate),
+ input: Arc::new(self.optimize(input)?),
+ }),
+ // Rest: recurse into plan, apply optimization where possible
+ LogicalPlan::Projection { .. }
+ | LogicalPlan::Aggregate { .. }
+ | LogicalPlan::Limit { .. }
+ | LogicalPlan::Repartition { .. }
+ | LogicalPlan::CreateExternalTable { .. }
+ | LogicalPlan::Extension { .. }
+ | LogicalPlan::Sort { .. }
+ | LogicalPlan::Explain { .. }
+ | LogicalPlan::Join { .. } => {
+ let expr = utils::expressions(plan);
+
+ // apply the optimization to all inputs of the plan
+ let inputs = utils::inputs(plan);
+ let new_inputs = inputs
+ .iter()
+ .map(|plan| self.optimize(plan))
+ .collect::<Result<Vec<_>>>()?;
+
+ utils::from_plan(plan, &expr, &new_inputs)
+ }
+ LogicalPlan::TableScan { .. } | LogicalPlan::EmptyRelation { .. } => {
+ Ok(plan.clone())
+ }
+ }
+ }
+
+ fn name(&self) -> &str {
+ "boolean_comparison"
+ }
+}
+
+/// Recursively transverses the logical plan.
+fn optimize_expr(e: &Expr) -> Expr {
+ match e {
+ Expr::BinaryExpr { left, op, right } => {
+ let left = optimize_expr(left);
+ let right = optimize_expr(right);
+ match op {
+ Operator::Eq => match (&left, &right) {
+ (Expr::Literal(ScalarValue::Boolean(b)), _) => match b {
+ Some(true) => right,
+ Some(false) | None => Expr::Not(Box::new(right)),
+ },
+ (_, Expr::Literal(ScalarValue::Boolean(b))) => match b {
+ Some(true) => left,
+ Some(false) | None => Expr::Not(Box::new(left)),
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: Operator::Eq,
+ right: Box::new(right),
+ },
+ },
+ Operator::NotEq => match (&left, &right) {
+ (Expr::Literal(ScalarValue::Boolean(b)), _) => match b {
+ Some(false) | None => right,
+ Some(true) => Expr::Not(Box::new(right)),
+ },
+ (_, Expr::Literal(ScalarValue::Boolean(b))) => match b {
+ Some(false) | None => left,
+ Some(true) => Expr::Not(Box::new(left)),
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: Operator::NotEq,
+ right: Box::new(right),
+ },
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: op.clone(),
+ right: Box::new(right),
+ },
+ }
+ }
+ Expr::Not(expr) => Expr::Not(Box::new(optimize_expr(&expr))),
+ Expr::Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ } => {
+ if expr.is_none() {
+ // recurse into CASE WHEN condition expressions
+ Expr::Case {
+ expr: None,
+ when_then_expr: when_then_expr
+ .iter()
+ .map(|(when, then)| (Box::new(optimize_expr(when)), then.clone()))
+ .collect(),
+ else_expr: else_expr.clone(),
+ }
+ } else {
+ // when base expression is specified, when_then_expr conditions are literal values
+ // so we can just skip this case
+ e.clone()
+ }
+ }
+ Expr::Alias { .. }
+ | Expr::Negative { .. }
+ | Expr::Column { .. }
+ | Expr::InList { .. }
+ | Expr::IsNotNull { .. }
+ | Expr::IsNull { .. }
+ | Expr::Cast { .. }
+ | Expr::ScalarVariable { .. }
+ | Expr::Between { .. }
+ | Expr::Literal { .. }
+ | Expr::ScalarFunction { .. }
+ | Expr::ScalarUDF { .. }
+ | Expr::AggregateFunction { .. }
+ | Expr::AggregateUDF { .. }
+ | Expr::Sort { .. }
+ | Expr::Wildcard => e.clone(),
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::logical_plan::{col, lit, LogicalPlanBuilder};
+ use crate::test::*;
+
+ fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
+ let mut rule = BooleanComparison::new();
+ let optimized_plan = rule.optimize(plan).expect("failed to optimize plan");
+ let formatted_plan = format!("{:?}", optimized_plan);
+ assert_eq!(formatted_plan, expected);
+ }
+
+ #[test]
+ fn simplify_eq_expr() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let plan = LogicalPlanBuilder::from(&table_scan)
+ .filter(col("a").eq(lit(true)))?
+ .filter(col("b").eq(lit(false)))?
+ .project(vec![col("a")])?
+ .build()?;
+
+ let expected = "\
+ Projection: #a\
+ \n Filter: NOT #b\
+ \n Filter: #a\
+ \n TableScan: test projection=None";
+
+ assert_optimized_plan_eq(&plan, expected);
+ Ok(())
+ }
+
+ #[test]
+ fn simplify_not_eq_expr() -> Result<()> {
Review comment:
One suggestion on the tests (not needed in this PR) would be to separtate out the cases --
1. one set in terms of `LogicalPlan` that makes sure rewrite was applied to all expressions in the plan node
2. Another set for just the logic in `optimize_expr` (aka make an `Expr`, rewrite it, and verify the rewrite was done correctly
##########
File path: rust/datafusion/src/optimizer/boolean_comparison.rs
##########
@@ -0,0 +1,270 @@
+// 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.
+
+//! Boolean comparision rule rewrites redudant comparison expression involing boolean literal into
+//! unary expression.
+
+use std::sync::Arc;
+
+use crate::error::Result;
+use crate::logical_plan::{Expr, LogicalPlan, Operator};
+use crate::optimizer::optimizer::OptimizerRule;
+use crate::optimizer::utils;
+use crate::scalar::ScalarValue;
+
+/// Optimizer that simplifies comparison expressions involving boolean literals.
+///
+/// Recursively go through all expressionss and simplify the following cases:
+/// * `expr = ture` to `expr`
+/// * `expr = false` to `!expr`
+pub struct BooleanComparison {}
+
+impl BooleanComparison {
Review comment:
As you can probably guess given PRs like https://github.com/apache/arrow/pull/9278 my preference to avoid repeating the structure walking logic is via a Visitor. Perhaps after this PR is merged, I can take a shot at rewriting it using a general `ExprRewriter` type pattern.
##########
File path: rust/datafusion/src/optimizer/boolean_comparison.rs
##########
@@ -0,0 +1,270 @@
+// 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.
+
+//! Boolean comparision rule rewrites redudant comparison expression involing boolean literal into
+//! unary expression.
+
+use std::sync::Arc;
+
+use crate::error::Result;
+use crate::logical_plan::{Expr, LogicalPlan, Operator};
+use crate::optimizer::optimizer::OptimizerRule;
+use crate::optimizer::utils;
+use crate::scalar::ScalarValue;
+
+/// Optimizer that simplifies comparison expressions involving boolean literals.
+///
+/// Recursively go through all expressionss and simplify the following cases:
+/// * `expr = ture` to `expr`
+/// * `expr = false` to `!expr`
+pub struct BooleanComparison {}
+
+impl BooleanComparison {
Review comment:
@Dandandan -- I think this PR as written is quite efficient and doesn't need a convergence loop as you suggest (which I think ends up potentially being quite inefficient if many rewrites are required) -- it already does a depth first traversal of the tree, simplifying on the way up.
I think convergence loops might be best used if we have several rules that can each potentially make changes that would unlock additional optimizations of the others
For example, if you had two different optimization functions like `optimization_A` and `optimization_B` but parts of `optimization_A` wouldn't be applied unless you ran `optimization_B`. In that case a loop like the following would let you take full advantage of that
```
while !changed {
let exprA = optimization_A(expr);
let exprB = optimization_B(exprA);
changed = expr != exprB;
expr = exprB
}
```
##########
File path: rust/datafusion/src/optimizer/boolean_comparison.rs
##########
@@ -0,0 +1,270 @@
+// 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.
+
+//! Boolean comparision rule rewrites redudant comparison expression involing boolean literal into
+//! unary expression.
+
+use std::sync::Arc;
+
+use crate::error::Result;
+use crate::logical_plan::{Expr, LogicalPlan, Operator};
+use crate::optimizer::optimizer::OptimizerRule;
+use crate::optimizer::utils;
+use crate::scalar::ScalarValue;
+
+/// Optimizer that simplifies comparison expressions involving boolean literals.
+///
+/// Recursively go through all expressionss and simplify the following cases:
+/// * `expr = ture` to `expr`
+/// * `expr = false` to `!expr`
+pub struct BooleanComparison {}
+
+impl BooleanComparison {
+ #[allow(missing_docs)]
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl OptimizerRule for BooleanComparison {
+ fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
+ match plan {
+ LogicalPlan::Filter { predicate, input } => Ok(LogicalPlan::Filter {
+ predicate: optimize_expr(predicate),
+ input: Arc::new(self.optimize(input)?),
+ }),
+ // Rest: recurse into plan, apply optimization where possible
+ LogicalPlan::Projection { .. }
+ | LogicalPlan::Aggregate { .. }
+ | LogicalPlan::Limit { .. }
+ | LogicalPlan::Repartition { .. }
+ | LogicalPlan::CreateExternalTable { .. }
+ | LogicalPlan::Extension { .. }
+ | LogicalPlan::Sort { .. }
+ | LogicalPlan::Explain { .. }
+ | LogicalPlan::Join { .. } => {
+ let expr = utils::expressions(plan);
+
+ // apply the optimization to all inputs of the plan
+ let inputs = utils::inputs(plan);
+ let new_inputs = inputs
+ .iter()
+ .map(|plan| self.optimize(plan))
+ .collect::<Result<Vec<_>>>()?;
+
+ utils::from_plan(plan, &expr, &new_inputs)
+ }
+ LogicalPlan::TableScan { .. } | LogicalPlan::EmptyRelation { .. } => {
+ Ok(plan.clone())
+ }
+ }
+ }
+
+ fn name(&self) -> &str {
+ "boolean_comparison"
+ }
+}
+
+/// Recursively transverses the logical plan.
+fn optimize_expr(e: &Expr) -> Expr {
+ match e {
+ Expr::BinaryExpr { left, op, right } => {
+ let left = optimize_expr(left);
+ let right = optimize_expr(right);
+ match op {
+ Operator::Eq => match (&left, &right) {
+ (Expr::Literal(ScalarValue::Boolean(b)), _) => match b {
+ Some(true) => right,
+ Some(false) | None => Expr::Not(Box::new(right)),
+ },
+ (_, Expr::Literal(ScalarValue::Boolean(b))) => match b {
+ Some(true) => left,
+ Some(false) | None => Expr::Not(Box::new(left)),
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: Operator::Eq,
+ right: Box::new(right),
+ },
+ },
+ Operator::NotEq => match (&left, &right) {
+ (Expr::Literal(ScalarValue::Boolean(b)), _) => match b {
+ Some(false) | None => right,
+ Some(true) => Expr::Not(Box::new(right)),
+ },
+ (_, Expr::Literal(ScalarValue::Boolean(b))) => match b {
+ Some(false) | None => left,
+ Some(true) => Expr::Not(Box::new(left)),
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: Operator::NotEq,
+ right: Box::new(right),
+ },
+ },
+ _ => Expr::BinaryExpr {
+ left: Box::new(left),
+ op: op.clone(),
+ right: Box::new(right),
+ },
+ }
+ }
+ Expr::Not(expr) => Expr::Not(Box::new(optimize_expr(&expr))),
+ Expr::Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ } => {
+ if expr.is_none() {
+ // recurse into CASE WHEN condition expressions
+ Expr::Case {
+ expr: None,
+ when_then_expr: when_then_expr
+ .iter()
+ .map(|(when, then)| (Box::new(optimize_expr(when)), then.clone()))
+ .collect(),
+ else_expr: else_expr.clone(),
Review comment:
I think you could also apply `optimize_expr` to then and else_expr:
```suggestion
.map(|(when, then)| (Box::new(optimize_expr(when)), optimize_expr(then))
.collect(),
else_expr: optimize_expr(else_expr),
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org