Fix wildcard keys for USING outer joins

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);

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,19 +443,86 @@ fn exclude_using_columns(plan: &LogicalPlan) -> Result<HashSet<Column>> {
     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<Column>,
+) -> Result<HashMap<Column, Expr>> {
+    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,
     plan: &LogicalPlan,
     wildcard_options: Option<&WildcardOptions>,
 ) -> Result<Vec<Expr>> {
     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.

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),