diff --git a/datafusion/optimizer/src/propagate_empty_relation.rs b/datafusion/optimizer/src/propagate_empty_relation.rs index bd96c558613a..f9183c467f62 100644 --- a/datafusion/optimizer/src/propagate_empty_relation.rs +++ b/datafusion/optimizer/src/propagate_empty_relation.rs @@ -210,11 +210,21 @@ impl OptimizerRule for PropagateEmptyRelation { if child.schema().eq(plan.schema()) { Ok(Transformed::yes(child)) } else { + // The union takes its field names from its first + // input, which may be gone. Map the remaining + // input's columns by position. + let exprs = child + .schema() + .columns() + .into_iter() + .zip(plan.schema().iter()) + .map(|(col, (qualifier, field))| { + Expr::Column(col) + .alias_qualified(qualifier.cloned(), field.name()) + }) + .collect::>(); Ok(Transformed::yes(LogicalPlan::Projection( - Projection::new_from_schema( - Arc::new(child), - Arc::clone(plan.schema()), - ), + Projection::try_new(exprs, Arc::new(child))?, ))) } } else { @@ -439,7 +449,8 @@ mod tests { let plan = LogicalPlanBuilder::from(left).union(right)?.build()?; - let expected = "Projection: a, b, c\n TableScan: test"; + let expected = + "Projection: test.a AS a, test.b AS b, test.c AS c\n TableScan: test"; assert_together_optimized_plan(plan, expected, true) } @@ -529,7 +540,34 @@ mod tests { let plan = LogicalPlanBuilder::from(left).union(right)?.build()?; - let expected = "Projection: a, b, c\n TableScan: test"; + let expected = + "Projection: test.a AS a, test.b AS b, test.c AS c\n TableScan: test"; + assert_together_optimized_plan(plan, expected, true) + } + + #[test] + fn propagate_union_remaining_input_with_other_names() -> Result<()> { + // Union + // Filter: false + // TableScan: test + // Projection: test2.a AS x, test2.b AS y, test2.c AS z + // TableScan: test2 + let left = LogicalPlanBuilder::from(test_table_scan()?) + .filter(lit(false))? + .build()?; + let right = LogicalPlanBuilder::from(test_table_scan_with_name("test2")?) + .project(vec![ + col("a").alias("x"), + col("b").alias("y"), + col("c").alias("z"), + ])? + .build()?; + + let plan = LogicalPlanBuilder::from(left).union(right)?.build()?; + + let expected = "Projection: x AS a, y AS b, z AS c\ + \n Projection: test2.a AS x, test2.b AS y, test2.c AS z\ + \n TableScan: test2"; assert_together_optimized_plan(plan, expected, true) } @@ -813,7 +851,7 @@ mod tests { .union(three)? .build()?; - let expected = "Projection: a, b, c\ + let expected = "Projection: test.a AS a, test.b AS b, test.c AS c\ \n TableScan: test"; assert_together_optimized_plan(plan, expected, true) diff --git a/datafusion/sqllogictest/test_files/struct.slt b/datafusion/sqllogictest/test_files/struct.slt index 6e0e1b483b1e..5b2ef30b200f 100644 --- a/datafusion/sqllogictest/test_files/struct.slt +++ b/datafusion/sqllogictest/test_files/struct.slt @@ -1982,8 +1982,8 @@ statement ok drop table rename_swap_struct_field; # Merging an extraction projection into a projection whose output lost its -# qualifier. Eliminating the empty side of the union leaves a projection of -# bare column names over a qualified input, and the merge used to add the +# qualifier. Eliminating the empty side of the union leaves a projection that +# renames qualified input columns to bare names, and the merge used to add the # pass-through columns under those bare names beside the qualified ones the # projection already carried, which is an ambiguous schema. statement ok @@ -2009,7 +2009,7 @@ logical_plan 03)----Aggregate: groupBy=[[expanded.env]], aggr=[[sum(CAST(__datafusion_extracted_1 AS Int64))]] 04)------SubqueryAlias: expanded 05)--------SubqueryAlias: samples -06)----------Projection: leaf_merge_source.env, get_field(s, Utf8("a")) AS __datafusion_extracted_1 +06)----------Projection: leaf_merge_source.env, get_field(leaf_merge_source.s, Utf8("a")) AS __datafusion_extracted_1 07)------------TableScan: leaf_merge_source projection=[s, env] statement ok diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index 2c706596b540..e5694327624c 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -1144,6 +1144,20 @@ query TITT ---- ns3 4 t3 m +# repro for https://github.com/apache/datafusion/issues/26516 +# should not error when only a later union input remains +query I +SELECT 1 WHERE false UNION ALL (SELECT 2 LIMIT 1); +---- +2 + +query I +SELECT a FROM (VALUES (1)) t(a) WHERE false +UNION ALL +(SELECT b FROM (VALUES (2)) s(b) LIMIT 1); +---- +2 + # Config reset statement ok RESET datafusion.execution.batch_size;