fix: resolve pass-through columns against the input when merging an extraction projection

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<Column> {
+    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<Column> = 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
     // =========================================================================

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<a int>, 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;