diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index 7c0e89ee96418..f1ffd6c5b12d0 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -25,6 +25,33 @@ use datafusion_sql::unparser::plan_to_sql; use super::*; +#[tokio::test] +async fn natural_full_join_wildcard_coalesces_join_key() -> Result<()> { + let ctx = SessionContext::new(); + let dataframe = ctx + .sql( + "WITH d1(id, name) AS (VALUES (1, 'a'), (2, 'b'), (4, 'c')), + d2(id, value) AS (VALUES (1, 'xx'), (2, 'yy'), (5, 'zz')) + SELECT * FROM d1 NATURAL FULL OUTER JOIN d2 ORDER BY id", + ) + .await?; + + assert_batches_eq!( + [ + "+----+------+-------+", + "| id | name | value |", + "+----+------+-------+", + "| 1 | a | xx |", + "| 2 | b | yy |", + "| 4 | c | |", + "| 5 | | zz |", + "+----+------+-------+", + ], + &dataframe.collect().await? + ); + Ok(()) +} + #[tokio::test] async fn join_change_in_planner() -> Result<()> { let config = SessionConfig::new().with_target_partitions(8); diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 297fbce14c795..84185ffecff72 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -24,7 +24,8 @@ use std::sync::Arc; use crate::expr::{Alias, Sort, WildcardOptions, WindowFunctionParams}; use crate::expr_rewriter::strip_outer_reference; use crate::{ - BinaryExpr, Expr, ExprSchemable, Filter, GroupingSet, LogicalPlan, Operator, and, + BinaryExpr, Expr, ExprSchemable, Filter, GroupingSet, JoinConstraint, JoinType, + LogicalPlan, Operator, and, when, }; use datafusion_expr_common::signature::{Signature, TypeSignature}; @@ -442,6 +443,63 @@ fn exclude_using_columns(plan: &LogicalPlan) -> Result> { Ok(excluded) } +/// Adjusts an unqualified wildcard over a top-level `USING` join so the +/// retained join key has the value required by SQL outer-join semantics. +fn using_join_wildcard_replacements( + plan: &LogicalPlan, + columns_to_skip: &mut HashSet, +) -> Result> { + let LogicalPlan::Join(join) = plan else { + return Ok(HashMap::new()); + }; + if join.join_constraint != JoinConstraint::Using { + return Ok(HashMap::new()); + } + if !matches!( + join.join_type, + JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full + ) { + return Ok(HashMap::new()); + } + + let mut replacements = HashMap::new(); + for (left_expr, right_expr) in &join.on { + let Some(left) = left_expr.get_as_join_column() else { + return internal_err!( + "Invalid USING join key. Expected column, found {left_expr:?}" + ); + }; + let Some(right) = right_expr.get_as_join_column() else { + return internal_err!( + "Invalid USING join key. Expected column, found {right_expr:?}" + ); + }; + + // Keep the left key in its original schema position and remove the + // duplicate right key. RIGHT and FULL joins replace its value below. + columns_to_skip.remove(left); + columns_to_skip.insert(right.clone()); + + let left_column = Expr::Column(left.clone()); + let right_column = Expr::Column(right.clone()); + let replacement = match join.join_type { + JoinType::Right => Some(right_column), + JoinType::Full => Some( + when(left_column.clone().is_not_null(), left_column) + .otherwise(right_column)?, + ), + _ => None, + }; + if let Some(replacement) = replacement { + replacements.insert( + left.clone(), + replacement.alias_qualified(left.relation.clone(), left.name.clone()), + ); + } + } + Ok(replacements) +} + /// Resolves an `Expr::Wildcard` to a collection of `Expr::Column`'s. pub fn expand_wildcard( schema: &DFSchema, @@ -449,12 +507,22 @@ pub fn expand_wildcard( wildcard_options: Option<&WildcardOptions>, ) -> Result> { let mut columns_to_skip = exclude_using_columns(plan)?; + let replacements = using_join_wildcard_replacements(plan, &mut columns_to_skip)?; columns_to_skip.extend(excluded_columns_from_schema( schema, wildcard_options, None, )?); - Ok(get_exprs_except_skipped(schema, &columns_to_skip)) + Ok(get_exprs_except_skipped(schema, &columns_to_skip) + .into_iter() + .map(|expr| match expr { + Expr::Column(column) => replacements + .get(&column) + .cloned() + .unwrap_or(Expr::Column(column)), + expr => expr, + }) + .collect()) } /// Resolves an unqualified wildcard using only the input schema. diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index d980a4de7eb21..8851b3212a6c3 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -5622,6 +5622,30 @@ fn test_using_join_wildcard_schema() { ] ); + // RIGHT and FULL joins must retain values from the non-null side of the + // merged USING column while preserving the wildcard schema. + let sql = "WITH t1 AS (SELECT 1 AS id, 'a' AS value1), + t2 AS (SELECT 2 AS id, 'x' AS value2) + SELECT * FROM t1 RIGHT JOIN t2 USING (id)"; + let plan = logical_plan(sql).unwrap(); + assert!( + plan.display_indent() + .to_string() + .contains("Projection: t2.id AS id, t1.value1, t2.value2"), + "{plan}" + ); + + let sql = "WITH t1 AS (SELECT 1 AS id, 'a' AS value1), + t2 AS (SELECT 2 AS id, 'x' AS value2) + SELECT * FROM t1 FULL OUTER JOIN t2 USING (id)"; + let plan = logical_plan(sql).unwrap(); + assert!( + plan.display_indent().to_string().contains( + "Projection: CASE WHEN t1.id IS NOT NULL THEN t1.id ELSE t2.id END AS id, t1.value1, t2.value2" + ), + "{plan}" + ); + // Multiple joins let sql = "WITH t1 AS (SELECT 1 AS a, 1 AS b), t2 AS (SELECT 1 AS a, 2 AS c),