Skip to content

Commit

Permalink
support cast/try_cast expr in reduceOuterJoin (#3621)
Browse files Browse the repository at this point in the history
* support cast/try_cast expr in reduceOuterJoin

* add new test case for cast/try_cast expr in reduceOuterJoin
  • Loading branch information
HuSen8891 authored Sep 29, 2022
1 parent 06a4f79 commit 9af2337
Showing 1 changed file with 41 additions and 1 deletion.
42 changes: 41 additions & 1 deletion datafusion/optimizer/src/reduce_outer_join.rs
Original file line number Diff line number Diff line change
Expand Up @@ -350,6 +350,15 @@ fn extract_nonnullable_columns(
false,
)
}
Expr::Cast { expr, data_type: _ } | Expr::TryCast { expr, data_type: _ } => {
extract_nonnullable_columns(
expr,
nonnullable_cols,
left_schema,
right_schema,
false,
)
}
_ => Ok(()),
}
}
Expand All @@ -358,9 +367,11 @@ fn extract_nonnullable_columns(
mod tests {
use super::*;
use crate::test::*;
use arrow::datatypes::DataType;
use datafusion_expr::{
binary_expr, col, lit,
binary_expr, cast, col, lit,
logical_plan::builder::LogicalPlanBuilder,
try_cast,
Operator::{And, Or},
};

Expand Down Expand Up @@ -481,4 +492,33 @@ mod tests {

Ok(())
}

#[test]
fn reduce_full_with_type_cast() -> Result<()> {
let t1 = test_table_scan_with_name("t1")?;
let t2 = test_table_scan_with_name("t2")?;

// reduce to inner join
let plan = LogicalPlanBuilder::from(t1)
.join(
&t2,
JoinType::Full,
(vec![Column::from_name("a")], vec![Column::from_name("a")]),
None,
)?
.filter(binary_expr(
cast(col("t1.b"), DataType::Int64).gt(lit(10u32)),
And,
try_cast(col("t2.c"), DataType::Int64).lt(lit(20u32)),
))?
.build()?;
let expected = "\
Filter: CAST(#t1.b AS Int64) > UInt32(10) AND TRY_CAST(#t2.c AS Int64) < UInt32(20)\
\n Inner Join: #t1.a = #t2.a\
\n TableScan: t1\
\n TableScan: t2";
assert_optimized_plan_eq(&plan, expected);

Ok(())
}
}

0 comments on commit 9af2337

Please sign in to comment.