Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use ExprRewriter in pre_cast_lit_in_comparison #3260

Merged
merged 2 commits into from
Aug 25, 2022
Merged

Use ExprRewriter in pre_cast_lit_in_comparison #3260

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bc3e967
Use ExprRewriter in pre_cast_lit_in_comparison.rs
andygrove Aug 25, 2022
6e0f926
remove manual recursion and add a nested test case
andygrove Aug 25, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 94 additions & 65 deletions datafusion/optimizer/src/pre_cast_lit_in_comparison.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
use crate::{OptimizerConfig, OptimizerRule};
use arrow::datatypes::DataType;
use datafusion_common::{DFSchemaRef, DataFusionError, Result, ScalarValue};
use datafusion_expr::expr_rewriter::{ExprRewritable, ExprRewriter, RewriteRecursion};
use datafusion_expr::utils::from_plan;
use datafusion_expr::{binary_expr, lit, Expr, ExprSchemable, LogicalPlan, Operator};

Expand Down Expand Up @@ -74,79 +75,93 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
.collect::<Result<Vec<_>>>()?;

let schema = plan.schema();

let mut expr_rewriter = PreCastLitExprRewriter {
schema: schema.clone(),
};

let new_exprs = plan
.expressions()
.into_iter()
.map(|expr| visit_expr(expr, schema))
.map(|expr| expr.rewrite(&mut expr_rewriter))
.collect::<Result<Vec<_>>>()?;

from_plan(plan, new_exprs.as_slice(), new_inputs.as_slice())
}

// Visit all type of expr, if the current has child expr, the child expr needed to visit first.
fn visit_expr(expr: Expr, schema: &DFSchemaRef) -> Result<Expr> {
// traverse the expr by dfs
match &expr {
Expr::BinaryExpr { left, op, right } => {
// dfs visit the left and right expr
let left = visit_expr(*left.clone(), schema)?;
let right = visit_expr(*right.clone(), schema)?;
let left_type = left.get_type(schema);
let right_type = right.get_type(schema);
// can't get the data type, just return the expr
if left_type.is_err() || right_type.is_err() {
return Ok(expr.clone());
}
let left_type = left_type.unwrap();
let right_type = right_type.unwrap();
if !left_type.eq(&right_type)
&& is_support_data_type(&left_type)
&& is_support_data_type(&right_type)
&& is_comparison_op(op)
{
match (&left, &right) {
(Expr::Literal(_), Expr::Literal(_)) => {
// do nothing
}
(Expr::Literal(left_lit_value), _)
if can_integer_literal_cast_to_type(
left_lit_value,
&right_type,
)? =>
{
// cast the left literal to the right type
return Ok(binary_expr(
cast_to_other_scalar_expr(left_lit_value, &right_type)?,
*op,
right,
));
}
(_, Expr::Literal(right_lit_value))
if can_integer_literal_cast_to_type(
right_lit_value,
&left_type,
)
.unwrap() =>
{
// cast the right literal to the left type
return Ok(binary_expr(
left,
*op,
cast_to_other_scalar_expr(right_lit_value, &left_type)?,
));
}
(_, _) => {
// do nothing
}
};
struct PreCastLitExprRewriter {
schema: DFSchemaRef,
}

impl ExprRewriter for PreCastLitExprRewriter {
fn pre_visit(&mut self, _expr: &Expr) -> Result<RewriteRecursion> {
Ok(RewriteRecursion::Continue)
}

fn mutate(&mut self, expr: Expr) -> Result<Expr> {
// traverse the expr by dfs
match &expr {
Expr::BinaryExpr { left, op, right } => {
// dfs visit the left and right expr
let left = self.mutate(*left.clone())?;
let right = self.mutate(*right.clone())?;
let left_type = left.get_type(&self.schema);
let right_type = right.get_type(&self.schema);
// can't get the data type, just return the expr
if left_type.is_err() || right_type.is_err() {
return Ok(expr.clone());
}
let left_type = left_type?;
let right_type = right_type?;
if !left_type.eq(&right_type)
&& is_support_data_type(&left_type)
&& is_support_data_type(&right_type)
&& is_comparison_op(op)
{
match (&left, &right) {
(Expr::Literal(_), Expr::Literal(_)) => {
// do nothing
}
(Expr::Literal(left_lit_value), _)
if can_integer_literal_cast_to_type(
left_lit_value,
&right_type,
)? =>
{
// cast the left literal to the right type
return Ok(binary_expr(
cast_to_other_scalar_expr(left_lit_value, &right_type)?,
*op,
right,
));
}
(_, Expr::Literal(right_lit_value))
if can_integer_literal_cast_to_type(
right_lit_value,
&left_type,
)
.unwrap() =>
{
// cast the right literal to the left type
return Ok(binary_expr(
left,
*op,
cast_to_other_scalar_expr(right_lit_value, &left_type)?,
));
}
(_, _) => {
// do nothing
}
};
}
// return the new binary op
Ok(binary_expr(left, *op, right))
}
// return the new binary op
Ok(binary_expr(left, *op, right))
// TODO: optimize in list
// Expr::InList { .. } => {}
// TODO: handle other expr type and dfs visit them
_ => Ok(expr),
}
// TODO: optimize in list
// Expr::InList { .. } => {}
// TODO: handle other expr type and dfs visit them
_ => Ok(expr),
}
}

Expand Down Expand Up @@ -245,9 +260,10 @@ fn can_integer_literal_cast_to_type(

#[cfg(test)]
mod tests {
use crate::pre_cast_lit_in_comparison::visit_expr;
use crate::pre_cast_lit_in_comparison::PreCastLitExprRewriter;
use arrow::datatypes::DataType;
use datafusion_common::{DFField, DFSchema, DFSchemaRef, ScalarValue};
use datafusion_expr::expr_rewriter::ExprRewritable;
use datafusion_expr::{col, lit, Expr};
use std::collections::HashMap;
use std::sync::Arc;
Expand Down Expand Up @@ -292,8 +308,21 @@ mod tests {
assert_eq!(optimize_test(c1_lt_lit_null, &schema), expected);
}

#[test]
fn nested() {
let schema = expr_test_schema();
// c1 < INT64(16) -> c1 < cast(INT32(16))
// the 16 is within the range of MAX(int32) and MIN(int32), we can cast the 16 to int32(16)
let expr_lt = col("c1").lt(lit(ScalarValue::Int64(Some(16)))).alias("x");
let expected = col("c1").lt(lit(ScalarValue::Int32(Some(16)))).alias("x");
assert_eq!(optimize_test(expr_lt, &schema), expected);
}

fn optimize_test(expr: Expr, schema: &DFSchemaRef) -> Expr {
visit_expr(expr, schema).unwrap()
let mut expr_rewriter = PreCastLitExprRewriter {
schema: schema.clone(),
};
expr.rewrite(&mut expr_rewriter).unwrap()
}

fn expr_test_schema() -> DFSchemaRef {
Expand Down