ygf11 commented on code in PR #4862: URL: https://github.com/apache/arrow-datafusion/pull/4862#discussion_r1067944748
########## datafusion/core/tests/sql/union.rs: ########## @@ -140,3 +140,85 @@ async fn union_schemas() -> Result<()> { assert_batches_eq!(expected, &result); Ok(()) } + +#[tokio::test] +async fn union_with_except_input() -> Result<()> { + let ctx = create_union_context()?; + let sql = "( + SELECT name FROM t1 + EXCEPT + SELECT name FROM t2 + ) + UNION ALL + ( + SELECT name FROM t2 + EXCEPT + SELECT name FROM t1 + )"; + let msg = format!("Creating logical plan for '{sql}'"); + let dataframe = ctx.sql(&("explain ".to_owned() + sql)).await.expect(&msg); + let plan = dataframe.into_optimized_plan()?; + let expected = vec![ + "Explain [plan_type:Utf8, plan:Utf8]", + " Union [name:UInt8;N]", + " LeftAnti Join: t1.name = t2.name [name:UInt8;N]", + " Distinct: [name:UInt8;N]", + " TableScan: t1 projection=[name] [name:UInt8;N]", + " Projection: t2.name [name:UInt8;N]", + " TableScan: t2 projection=[name] [name:UInt8;N]", + " LeftAnti Join: t2.name = t1.name [name:UInt8;N]", + " Distinct: [name:UInt8;N]", + " TableScan: t2 projection=[name] [name:UInt8;N]", + " Projection: t1.name [name:UInt8;N]", + " TableScan: t1 projection=[name] [name:UInt8;N]", + ]; + + let formatted = plan.display_indent_schema().to_string(); + let actual: Vec<&str> = formatted.trim().lines().collect(); + assert_eq!( + expected, actual, + "\n\nexpected:\n\n{expected:#?}\nactual:\n\n{actual:#?}\n\n" + ); + Ok(()) +} + +#[tokio::test] +async fn union_with_type_coercion() -> Result<()> { + let ctx = create_union_context()?; + let sql = "( + SELECT id, name FROM t1 + EXCEPT + SELECT id, name FROM t2 + ) + UNION ALL + ( + SELECT id, name FROM t2 + EXCEPT + SELECT id, name FROM t1 + )"; + let msg = format!("Creating logical plan for '{sql}'"); + let dataframe = ctx.sql(&("explain ".to_owned() + sql)).await.expect(&msg); + let plan = dataframe.into_optimized_plan()?; + let expected = vec![ + "Explain [plan_type:Utf8, plan:Utf8]", + " Union [id:Int32;N, name:UInt8;N]", + " LeftAnti Join: t1.id = CAST(t2.id AS Int32), t1.name = t2.name [id:Int32;N, name:UInt8;N]", + " Distinct: [id:Int32;N, name:UInt8;N]", + " TableScan: t1 projection=[id, name] [id:Int32;N, name:UInt8;N]", + " Projection: t2.id, t2.name [id:UInt8;N, name:UInt8;N]", + " TableScan: t2 projection=[id, name] [id:UInt8;N, name:UInt8;N]", + " Projection: CAST(t2.id AS Int32) AS id, t2.name [id:Int32;N, name:UInt8;N]", + " LeftAnti Join: CAST(t2.id AS Int32) = t1.id, t2.name = t1.name [id:UInt8;N, name:UInt8;N]", + " Distinct: [id:UInt8;N, name:UInt8;N]", + " TableScan: t2 projection=[id, name] [id:UInt8;N, name:UInt8;N]", + " Projection: t1.id, t1.name [id:Int32;N, name:UInt8;N]", + " TableScan: t1 projection=[id, name] [id:Int32;N, name:UInt8;N]", + ]; Review Comment: Same as the above test. -- 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: github-unsubscr...@arrow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org