in datafusion/optimizer/src/unwrap_cast_in_comparison.rs [133:269]
fn mutate(&mut self, expr: Expr) -> Result<Expr> {
match &expr {
// For case:
// try_cast/cast(expr as data_type) op literal
// literal op try_cast/cast(expr as data_type)
Expr::BinaryExpr(BinaryExpr { left, op, right }) => {
let left = left.as_ref().clone();
let right = right.as_ref().clone();
let left_type = left.get_type(&self.schema)?;
let right_type = right.get_type(&self.schema)?;
// Because the plan has been done the type coercion, the left and right must be equal
if is_support_data_type(&left_type)
&& is_support_data_type(&right_type)
&& is_comparison_op(op)
{
match (&left, &right) {
(
Expr::Literal(left_lit_value),
Expr::TryCast(TryCast { expr, .. })
| Expr::Cast(Cast { expr, .. }),
) => {
// if the left_lit_value can be casted to the type of expr
// we need to unwrap the cast for cast/try_cast expr, and add cast to the literal
let expr_type = expr.get_type(&self.schema)?;
let casted_scalar_value =
try_cast_literal_to_type(left_lit_value, &expr_type)?;
if let Some(value) = casted_scalar_value {
// unwrap the cast/try_cast for the right expr
return Ok(binary_expr(
lit(value),
*op,
expr.as_ref().clone(),
));
}
}
(
Expr::TryCast(TryCast { expr, .. })
| Expr::Cast(Cast { expr, .. }),
Expr::Literal(right_lit_value),
) => {
// if the right_lit_value can be casted to the type of expr
// we need to unwrap the cast for cast/try_cast expr, and add cast to the literal
let expr_type = expr.get_type(&self.schema)?;
let casted_scalar_value =
try_cast_literal_to_type(right_lit_value, &expr_type)?;
if let Some(value) = casted_scalar_value {
// unwrap the cast/try_cast for the left expr
return Ok(binary_expr(
expr.as_ref().clone(),
*op,
lit(value),
));
}
}
(_, _) => {
// do nothing
}
};
}
// return the new binary op
Ok(binary_expr(left, *op, right))
}
// For case:
// try_cast/cast(expr as left_type) in (expr1,expr2,expr3)
Expr::InList(InList {
expr: left_expr,
list,
negated,
}) => {
if let Some(
Expr::TryCast(TryCast {
expr: internal_left_expr,
..
})
| Expr::Cast(Cast {
expr: internal_left_expr,
..
}),
) = Some(left_expr.as_ref())
{
let internal_left = internal_left_expr.as_ref().clone();
let internal_left_type = internal_left.get_type(&self.schema);
if internal_left_type.is_err() {
// error data type
return Ok(expr);
}
let internal_left_type = internal_left_type?;
if !is_support_data_type(&internal_left_type) {
// not supported data type
return Ok(expr);
}
let right_exprs = list
.iter()
.map(|right| {
let right_type = right.get_type(&self.schema)?;
if !is_support_data_type(&right_type) {
return Err(DataFusionError::Internal(format!(
"The type of list expr {} not support",
&right_type
)));
}
match right {
Expr::Literal(right_lit_value) => {
// if the right_lit_value can be casted to the type of internal_left_expr
// we need to unwrap the cast for cast/try_cast expr, and add cast to the literal
let casted_scalar_value =
try_cast_literal_to_type(right_lit_value, &internal_left_type)?;
if let Some(value) = casted_scalar_value {
Ok(lit(value))
} else {
Err(DataFusionError::Internal(format!(
"Can't cast the list expr {:?} to type {:?}",
right_lit_value, &internal_left_type
)))
}
}
other_expr => Err(DataFusionError::Internal(format!(
"Only support literal expr to optimize, but the expr is {:?}",
&other_expr
))),
}
})
.collect::<Result<Vec<_>>>();
match right_exprs {
Ok(right_exprs) => {
Ok(in_list(internal_left, right_exprs, *negated))
}
Err(_) => Ok(expr),
}
} else {
Ok(expr)
}
}
// TODO: handle other expr type and dfs visit them
_ => Ok(expr),
}
}