ygf11 commented on code in PR #4353:
URL: https://github.com/apache/arrow-datafusion/pull/4353#discussion_r1031174418
##########
datafusion/optimizer/src/type_coercion.rs:
##########
@@ -1114,4 +1293,145 @@ mod test {
Ok(())
// TODO add more test for this
}
+
+ #[test]
+ fn test_column_join_with_different_data_types() -> Result<()> {
+ let t1 = join_test_table_scan_with_name("t1")?;
+ let t2 = join_test_table_scan_with_name("t2")?;
+
+ let join_plan = create_test_join_plan(
+ t1,
+ t2,
+ JoinType::Inner,
+ (vec![col("a")], vec![col("b")]),
+ None,
+ )?;
+
+ let expected = "Inner Join: t1.a#0 = t2.b\
+ \n Projection: t1.a, t1.b, t1.c, CAST(t1.a AS Int32) AS t1.a#0\
+ \n TableScan: t1\
+ \n TableScan: t2";
+ let rule = TypeCoercion::new();
+ let mut config = OptimizerConfig::default();
+ let plan = rule.optimize(&join_plan, &mut config).unwrap();
+ assert_eq!(&format!("{:?}", plan), expected);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_column_join_with_filter_and_type_coercion() -> Result<()> {
+ let t1 = join_test_table_scan_with_name("t1")?;
+ let t2 = join_test_table_scan_with_name("t2")?;
+
+ let join_plan = create_test_join_plan(
+ t1,
+ t2,
+ JoinType::Inner,
+ (vec![col("a")], vec![col("c")]),
+ Some(Expr::Column(Column::new("t1".into(), "a")).lt(lit(10i32))),
+ )?;
+
+ let expected =
+ "Inner Join: t1.a#0 = t2.c Filter: CAST(t1.a AS Int32) < Int32(10)\
+ \n Projection: t1.a, t1.b, t1.c, CAST(t1.a AS Int64) AS t1.a#0\
+ \n TableScan: t1\
+ \n TableScan: t2";
+ let rule = TypeCoercion::new();
+ let mut config = OptimizerConfig::default();
+ let plan = rule.optimize(&join_plan, &mut config).unwrap();
+ assert_eq!(&format!("{:?}", plan), expected);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_expr_join_with_different_data_types() -> Result<()> {
+ let t1 = join_test_table_scan_with_name("t1")?;
+ let t2 = join_test_table_scan_with_name("t2")?;
+
+ let left_keys = vec![col("a") + lit(1)];
+ let right_keys = vec![col("c") * lit(2)];
+ let join_plan = create_test_join_plan(
+ t1,
+ t2,
+ JoinType::Inner,
+ (left_keys, right_keys),
+ None,
+ )?;
+
+ let expected = "Projection: t1.a, t1.b, t1.c, t2.a, t2.b, t2.c\
+ \n Inner Join: t1.a + Int32(1)#0 = t2.c * Int32(2)\
+ \n Projection: t1.a, t1.b, t1.c, t1.a + Int32(1), CAST(t1.a +
Int32(1) AS Int64) AS t1.a + Int32(1)#0\
+ \n Projection: t1.a, t1.b, t1.c, CAST(t1.a AS Int32) + Int32(1)\
+ \n TableScan: t1\
+ \n Projection: t2.a, t2.b, t2.c, t2.c * CAST(Int32(2) AS Int64)\
+ \n TableScan: t2";
+ let rule = TypeCoercion::new();
+ let mut config = OptimizerConfig::default();
+ let plan = rule.optimize(&join_plan, &mut config).unwrap();
+ assert_eq!(&format!("{:?}", plan), expected);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_multiple_keys_join_type_coercion() -> Result<()> {
+ let t1 = join_test_table_scan_with_name("t1")?;
+ let t2 = join_test_table_scan_with_name("t2")?;
+
+ let join_plan = create_test_join_plan(
+ t1,
+ t2,
+ JoinType::Inner,
+ (
+ vec![col("a"), col("b"), col("c")],
+ vec![col("b"), col("c"), col("a")],
+ ),
+ None,
+ )?;
+
+ let expected = "Inner Join: t1.a#0 = t2.b, t1.b#1 = t2.c, t1.c =
t2.a#2\
+ \n Projection: t1.a, t1.b, t1.c, CAST(t1.a AS Int32) AS t1.a#0,
CAST(t1.b AS Int64) AS t1.b#1\
+ \n TableScan: t1\
+ \n Projection: t2.a, t2.b, t2.c, CAST(t2.a AS Int64) AS t2.a#2\
+ \n TableScan: t2";
+ let rule = TypeCoercion::new();
+ let mut config = OptimizerConfig::default();
+ let plan = rule.optimize(&join_plan, &mut config).unwrap();
+ assert_eq!(&format!("{:?}", plan), expected);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_same_key_multiple_time_join() -> Result<()> {
+ let t1 = join_test_table_scan_with_name("t1")?;
+ let t2 = join_test_table_scan_with_name("t2")?;
+
+ // a + 2 happen twice, and both need insert type coercion.
+ let join_plan = create_test_join_plan(
+ t1,
+ t2,
+ JoinType::Inner,
+ (
+ vec![col("a") * lit(2u8), col("a") * lit(2u8)],
+ vec![col("b"), col("c")],
+ ),
+ None,
+ )?;
+
+ let expected = "Projection: t1.a, t1.b, t1.c, t2.a, t2.b, t2.c\
+ \n Inner Join: t1.a * UInt8(2)#0 = t2.b, t1.a * UInt8(2)#1 = t2.c\
+ \n Projection: t1.a, t1.b, t1.c, t1.a * UInt8(2), CAST(t1.a *
UInt8(2) AS Int32) AS t1.a * UInt8(2)#0, CAST(t1.a * UInt8(2) AS Int64) AS t1.a
* UInt8(2)#1\
+ \n Projection: t1.a, t1.b, t1.c, t1.a * CAST(UInt8(2) AS UInt16)\
+ \n TableScan: t1\
+ \n TableScan: t2";
+ let rule = TypeCoercion::new();
Review Comment:
In this test case, the column happens twice and both need type coercion, I
give them an alias to distinguish from each other.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
