diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index 8a4abfcb48e56..e267a9d5b8751 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -637,6 +637,24 @@ impl<'a> LeafExpressionExtractor<'a> { } } +/// The way `schema` names `col`, or `None` when it does not hold it unambiguously. +/// +/// A qualified column is taken as it stands. An unqualified one is matched on name alone +/// and comes back carrying the qualifier the schema gives that field, so a column pushed +/// into a projection reads as the input spells it. A name the schema holds more than once +/// resolves to nothing rather than to an arbitrary one of them. +fn resolve_against(schema: &DFSchema, col: &Column) -> Option { + match &col.relation { + Some(relation) => schema + .has_column_with_qualified_name(relation, &col.name) + .then(|| col.clone()), + None => schema + .qualified_field_with_unqualified_name(&col.name) + .ok() + .map(|(qualifier, field)| Column::new(qualifier.cloned(), field.name())), + } +} + /// Build an extraction projection above the target node (shared by both passes). /// /// If the target is an existing projection, merges into it. This requires @@ -702,27 +720,31 @@ fn build_extraction_projection_impl( // than target_schema (the projection's output) because columns produced // by alias expressions (e.g., CSE's __common_expr_N) exist in the output but // not the input, and cannot be added as pass-through Column references. + // + // Both sides of that check are read in the input's own spelling. A column can + // arrive here unqualified while the input names it `t.c`, and the other way round: + // a projection of bare names over a qualified input is what eliminating one side + // of a union leaves behind. Resolving both sides lets a pass-through the + // projection already carries match the one about to be added, and pushes the one + // that is genuinely new under the name the input gives it. Left unresolved, the + // merged projection would hold `t.c` and a bare `c` together, which + // `Projection::try_new` rejects as ambiguous. A same-name alias such as + // `t.c AS c` is a pass-through as well, and counts the same as a bare `t.c`. + let input_schema = existing.input.schema(); let existing_cols: IndexSet = existing .expr .iter() - .filter_map(|e| { - if let Expr::Column(c) = e { - Some(c.clone()) - } else { - None - } - }) + .filter_map(|e| resolve_against(input_schema, passthrough_column(e)?)) .collect(); - let input_schema = existing.input.schema(); for col in columns_needed { let col_expr = Expr::Column(col.clone()); let resolved = replace_cols_by_name(col_expr, &replace_map)?; if let Expr::Column(resolved_col) = &resolved - && !existing_cols.contains(resolved_col) - && input_schema.has_column(resolved_col) + && let Some(input_col) = resolve_against(input_schema, resolved_col) + && !existing_cols.contains(&input_col) { - proj_exprs.push(Expr::Column(resolved_col.clone())); + proj_exprs.push(Expr::Column(input_col)); } // If resolved to non-column expr, it's already computed by existing projection } @@ -2350,6 +2372,83 @@ mod tests { "#) } + /// A filter can name a column bare (`a`) while the projection below it + /// outputs the qualified `test.a`, as eliminating the empty side of a union + /// leaves behind. The merge must match the two, and not add a bare `a` + /// beside `test.a`, which is an ambiguous schema. + #[test] + fn test_merge_bare_column_into_qualified_projection() -> Result<()> { + let table_scan = test_table_scan()?; + let projection = LogicalPlanBuilder::from(table_scan) + .project(vec![ + col("test.a"), + col("test.b"), + (col("test.c") + lit(1)).alias("d"), + ])? + .build()?; + let predicate = + leaf_udf(Expr::Column(Column::new_unqualified("a")), "x").eq(lit(1)); + let plan = LogicalPlan::Filter(datafusion_expr::Filter::try_new( + predicate, + Arc::new(projection), + )?); + + assert_stages!(plan, @r#" + ## Original Plan + Filter: leaf_udf(a, Utf8("x")) = Int32(1) + Projection: test.a, test.b, test.c + Int32(1) AS d + TableScan: test projection=[a, b, c] + + ## After Extraction + Projection: test.a, test.b, d + Filter: __datafusion_extracted_1 = Int32(1) + Projection: test.a, test.b, test.c + Int32(1) AS d, leaf_udf(a, Utf8("x")) AS __datafusion_extracted_1 + TableScan: test projection=[a, b, c] + + ## After Pushdown + (same as after extraction) + + ## Optimized + (same as after pushdown) + "#) + } + + /// A projection can spell a pass-through column as a same-name alias + /// (`test.a AS a`). Merging an extraction into it must treat that alias as + /// the pass-through it is, and not add `test.a` beside the `a` it outputs, + /// which is an ambiguous schema. + #[test] + fn test_merge_into_projection_with_same_name_alias() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .project(vec![ + col("test.a").alias("a"), + col("test.b").alias("b"), + col("test.c").alias("c"), + ])? + .filter(leaf_udf(col("a"), "x").eq(lit(1)))? + .build()?; + + assert_stages!(plan, @r#" + ## Original Plan + Filter: leaf_udf(a, Utf8("x")) = Int32(1) + Projection: test.a AS a, test.b AS b, test.c AS c + TableScan: test projection=[a, b, c] + + ## After Extraction + Projection: a, b, c + Filter: __datafusion_extracted_1 = Int32(1) + Projection: test.a AS a, test.b AS b, test.c AS c, leaf_udf(test.a, Utf8("x")) AS __datafusion_extracted_1 + TableScan: test projection=[a, b, c] + + ## After Pushdown + (same as after extraction) + + ## Optimized + (same as after pushdown) + "#) + } + // ========================================================================= // Join extraction tests // ========================================================================= diff --git a/datafusion/sqllogictest/test_files/struct.slt b/datafusion/sqllogictest/test_files/struct.slt index 87bbd11c986a4..10ecbf0885211 100644 --- a/datafusion/sqllogictest/test_files/struct.slt +++ b/datafusion/sqllogictest/test_files/struct.slt @@ -1803,3 +1803,28 @@ drop view struct_ctor_view; statement ok drop table struct_ctor_null; + +# 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 +# pass-through columns under those bare names beside the qualified ones the +# projection already carried, which is an ambiguous schema. +statement ok +create table leaf_merge_source(v int, s struct, env varchar) as values (1, {a: 10}, 'prod'), (2, {a: 20}, 'dev'); + +query TI +with samples as ( + select v, s, env from leaf_merge_source +), +expanded as ( + select v, s, env from samples + union all + select v, s, env from samples where 1 = 2 +) +select env, sum(s['a']) from expanded group by env order by env; +---- +dev 20 +prod 10 + +statement ok +drop table leaf_merge_source;