This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new d314ced809 Coerce types for all union children plans when eliminating
nesting (#11386)
d314ced809 is described below
commit d314ced8090cb599fd7808d7df41699e46ac956e
Author: Marko Grujic <[email protected]>
AuthorDate: Thu Jul 11 18:22:20 2024 +0200
Coerce types for all union children plans when eliminating nesting (#11386)
---
datafusion/optimizer/src/eliminate_nested_union.rs | 13 +++++++------
datafusion/sqllogictest/test_files/union.slt | 15 +++++++++++++++
2 files changed, 22 insertions(+), 6 deletions(-)
diff --git a/datafusion/optimizer/src/eliminate_nested_union.rs
b/datafusion/optimizer/src/eliminate_nested_union.rs
index c8ae937e12..cc8cf1f56c 100644
--- a/datafusion/optimizer/src/eliminate_nested_union.rs
+++ b/datafusion/optimizer/src/eliminate_nested_union.rs
@@ -60,7 +60,8 @@ impl OptimizerRule for EliminateNestedUnion {
let inputs = inputs
.into_iter()
.flat_map(extract_plans_from_union)
- .collect::<Vec<_>>();
+ .map(|plan| coerce_plan_expr_for_schema(&plan, &schema))
+ .collect::<Result<Vec<_>>>()?;
Ok(Transformed::yes(LogicalPlan::Union(Union {
inputs: inputs.into_iter().map(Arc::new).collect_vec(),
@@ -74,7 +75,8 @@ impl OptimizerRule for EliminateNestedUnion {
.into_iter()
.map(extract_plan_from_distinct)
.flat_map(extract_plans_from_union)
- .collect::<Vec<_>>();
+ .map(|plan| coerce_plan_expr_for_schema(&plan,
&schema))
+ .collect::<Result<Vec<_>>>()?;
Ok(Transformed::yes(LogicalPlan::Distinct(Distinct::All(
Arc::new(LogicalPlan::Union(Union {
@@ -95,10 +97,9 @@ impl OptimizerRule for EliminateNestedUnion {
fn extract_plans_from_union(plan: Arc<LogicalPlan>) -> Vec<LogicalPlan> {
match unwrap_arc(plan) {
- LogicalPlan::Union(Union { inputs, schema }) => inputs
- .into_iter()
- .map(|plan| coerce_plan_expr_for_schema(&plan, &schema).unwrap())
- .collect::<Vec<_>>(),
+ LogicalPlan::Union(Union { inputs, .. }) => {
+ inputs.into_iter().map(unwrap_arc).collect::<Vec<_>>()
+ }
plan => vec![plan],
}
}
diff --git a/datafusion/sqllogictest/test_files/union.slt
b/datafusion/sqllogictest/test_files/union.slt
index 7b91e97e4a..5ede68a42a 100644
--- a/datafusion/sqllogictest/test_files/union.slt
+++ b/datafusion/sqllogictest/test_files/union.slt
@@ -135,6 +135,21 @@ SELECT SUM(d) FROM (
----
5
+# three way union with aggregate and type coercion
+query II rowsort
+SELECT c1, SUM(c2) FROM (
+ SELECT 1 as c1, 1::int as c2
+ UNION
+ SELECT 2 as c1, 2::int as c2
+ UNION
+ SELECT 3 as c1, COALESCE(3::int, 0) as c2
+) as a
+GROUP BY c1
+----
+1 1
+2 2
+3 3
+
# union_all_with_count
statement ok
CREATE table t as SELECT 1 as a
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
