fix: hide internal USING merged-key qualifier

Keep wildcard output schemas on their visible input columns while still
evaluating RIGHT and FULL USING keys with merged-key semantics. Also
unwrap the reserved merged-key alias for GROUP BY planning so invalid
ORDER BY keys report the aggregate validation error instead of a schema
lookup failure.

Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>

Author: Phoenix500526

datafusion/expr/src/utils.rs

@@ -21,8 +21,10 @@ use std::cmp::Ordering;
 use std::collections::{BTreeSet, HashSet};
 use std::sync::Arc;
 
+use crate::conditional_expressions::CaseBuilder;
 use crate::expr::{Alias, Sort, WildcardOptions, WindowFunctionParams};
-use crate::expr_rewriter::{merged_using_key_or_column, strip_outer_reference};
+use crate::expr_rewriter::strip_outer_reference;
+use crate::logical_plan::JoinType;
 use crate::{
     BinaryExpr, Expr, ExprSchemable, Filter, GroupingSet, LogicalPlan, Operator, and,
 };
@@ -442,6 +444,38 @@ fn exclude_using_columns(plan: &LogicalPlan) -> Result<HashSet<Column>> {
     Ok(excluded)
 }
 
+fn merged_using_key_for_wildcard(
+    col: Column,
+    merged_keys: &[(Column, Column, JoinType)],
+) -> Result<Expr> {
+    let Some((l, r, join_type)) =
+        merged_keys.iter().find(|(l, r, _)| l == &col || r == &col)
+    else {
+        return Ok(Expr::Column(col));
+    };
+
+    let (expr, visible_col) = match join_type {
+        JoinType::Right | JoinType::RightSemi | JoinType::RightAnti => {
+            (Expr::Column(r.clone()), r.clone())
+        }
+        JoinType::Full => CaseBuilder::new(
+            None,
+            vec![Expr::Column(l.clone()).is_not_null()],
+            vec![Expr::Column(l.clone())],
+            Some(Box::new(Expr::Column(r.clone()))),
+        )
+        .end()
+        .map(|expr| (expr, col.clone()))?,
+        _ => (Expr::Column(l.clone()), l.clone()),
+    };
+
+    if expr == Expr::Column(visible_col.clone()) {
+        Ok(expr)
+    } else {
+        Ok(expr.alias_qualified(visible_col.relation, visible_col.name))
+    }
+}
+
 /// Resolves an `Expr::Wildcard` to a collection of `Expr::Column`'s.
 pub fn expand_wildcard(
     schema: &DFSchema,
@@ -463,16 +497,16 @@ pub fn expand_wildcard(
     columns_to_skip.extend(excluded_columns);
     let exprs = get_exprs_except_skipped(schema, &columns_to_skip);
 
-    // Resolve the surviving USING / NATURAL key column to the internally named
-    // merged key expression.
+    // Resolve the surviving USING / NATURAL key column to the merged key value,
+    // while preserving the wildcard-visible field qualifier.
     let merged_keys = plan.using_key_pairs()?;
     if merged_keys.is_empty() {
         return Ok(exprs);
     }
     exprs
         .into_iter()
         .map(|expr| match expr {
-            Expr::Column(col) => merged_using_key_or_column(col, true, &merged_keys),
+            Expr::Column(col) => merged_using_key_for_wildcard(col, &merged_keys),
             other => Ok(other),
         })
         .collect()

datafusion/sql/src/select.rs

@@ -40,7 +40,7 @@ use datafusion_common::{
 use datafusion_common::{RecursionUnnestOption, UnnestOptions};
 use datafusion_expr::ExprSchemable;
 use datafusion_expr::builder::get_struct_unnested_columns;
-use datafusion_expr::expr::{PlannedReplaceSelectItem, Unnest, WildcardOptions};
+use datafusion_expr::expr::{Alias, PlannedReplaceSelectItem, Unnest, WildcardOptions};
 use datafusion_expr::expr_rewriter::{
     merged_using_key_or_column, normalize_col, normalize_sorts,
 };
@@ -135,6 +135,40 @@ fn normalize_col_resolving_merged_using_key(
     .data()
 }
 
+fn unalias_internal_merged_key(expr: Expr, select_exprs: &[Expr]) -> Expr {
+    match expr {
+        Expr::Alias(Alias {
+            expr,
+            relation: Some(relation),
+            ..
+        }) if relation.table() == LogicalPlanBuilder::MERGED_KEY_QUALIFIER => *expr,
+        Expr::Column(column)
+            if column.relation.as_ref().is_some_and(|relation| {
+                relation.table() == LogicalPlanBuilder::MERGED_KEY_QUALIFIER
+            }) =>
+        {
+            select_exprs
+                .iter()
+                .find_map(|select_expr| match select_expr {
+                    Expr::Alias(Alias {
+                        expr,
+                        relation: Some(relation),
+                        name,
+                        ..
+                    }) if relation.table()
+                        == LogicalPlanBuilder::MERGED_KEY_QUALIFIER
+                        && name == &column.name =>
+                    {
+                        Some(*expr.clone())
+                    }
+                    _ => None,
+                })
+                .unwrap_or(Expr::Column(column))
+        }
+        other => other,
+    }
+}
+
 impl<S: ContextProvider> SqlToRel<'_, S> {
     /// Generate a logic plan from an SQL select
     pub(super) fn select_to_plan(
@@ -288,6 +322,8 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
                     let group_by_expr =
                         resolve_positions_to_exprs(group_by_expr, &select_exprs)?;
                     let group_by_expr = normalize_col(group_by_expr, &projected_plan)?;
+                    let group_by_expr =
+                        unalias_internal_merged_key(group_by_expr, &select_exprs);
                     self.validate_schema_satisfies_exprs(
                         base_plan.schema(),
                         std::slice::from_ref(&group_by_expr),

datafusion/sql/tests/sql_integration.rs

@@ -1074,7 +1074,7 @@ fn join_with_ambiguous_column() {
     assert_snapshot!(
         plan,
         @r"
-    Projection: a.id
+    Projection: a.id AS id
       Inner Join: Using a.id = b.id
         SubqueryAlias: a
           TableScan: person
@@ -1091,7 +1091,7 @@ fn natural_left_join() {
     assert_snapshot!(
         plan,
         @r"
-    Projection: a.l_item_id
+    Projection: a.l_item_id AS l_item_id
       Left Join: Using a.l_orderkey = b.l_orderkey, a.l_item_id = b.l_item_id, a.l_description = b.l_description, a.l_extendedprice = b.l_extendedprice, a.price = b.price
         SubqueryAlias: a
           TableScan: lineitem
@@ -1108,7 +1108,7 @@ fn natural_right_join() {
     assert_snapshot!(
         plan,
         @r"
-    Projection: b.l_item_id
+    Projection: b.l_item_id AS l_item_id
       Right Join: Using a.l_orderkey = b.l_orderkey, a.l_item_id = b.l_item_id, a.l_description = b.l_description, a.l_extendedprice = b.l_extendedprice, a.price = b.price
         SubqueryAlias: a
           TableScan: lineitem
@@ -2551,7 +2551,7 @@ fn join_with_using() {
     assert_snapshot!(
         plan,
         @r"
-    Projection: person.first_name, person.id
+    Projection: person.first_name, person.id AS id
       Inner Join: Using person.id = person2.id
         TableScan: person
         SubqueryAlias: person2
@@ -3502,6 +3502,17 @@ fn select_groupby_orderby_aggregate_on_non_selected_column_original_issue() {
     );
 }
 
+#[test]
+fn natural_join_group_by_order_by_non_grouped_merged_key() {
+    let sql = "WITH t1 AS (SELECT 1 AS v1, 2 AS v2) \
+        SELECT v1 FROM t1 AS tt1 NATURAL JOIN t1 AS tt2 \
+        GROUP BY v1 ORDER BY v2";
+    assert_contains!(
+        error_message(sql),
+        "Column in ORDER BY must be in GROUP BY or an aggregate function"
+    );
+}
+
 fn logical_plan(sql: &str) -> Result<LogicalPlan> {
     logical_plan_with_options(sql, ParserOptions::default())
 }

datafusion/sqllogictest/test_files/join.slt.part

@@ -1167,7 +1167,7 @@ INNER JOIN t0 ON (t0.v1 + t5.v0) > 0
 WHERE t0.v1 = t1.v0;
 ----
 logical_plan
-01)Projection: t1.v0 AS v0, t1.v1 AS v1, t5.v2, t5.v3, t5.v4, t0.v0, t0.v1
+01)Projection: t1.v0, t1.v1, t5.v2, t5.v3, t5.v4, t0.v0, t0.v1
 02)--Inner Join: CAST(t1.v0 AS Float64) = t0.v1 Filter: t0.v1 + CAST(t5.v0 AS Float64) > Float64(0)
 03)----Projection: t1.v0, t1.v1, t5.v0, t5.v2, t5.v3, t5.v4
 04)------Inner Join: t1.v0 = t5.v0, t1.v1 = t5.v1

datafusion/sqllogictest/test_files/joins.slt

@@ -4542,7 +4542,7 @@ query TT
 explain SELECT * FROM person a join person b using (id, age);
 ----
 logical_plan
-01)Projection: a.id AS id, a.age AS age, a.state, b.state
+01)Projection: a.id, a.age, a.state, b.state
 02)--Inner Join: a.id = b.id, a.age = b.age
 03)----SubqueryAlias: a
 04)------TableScan: person projection=[id, age, state]
@@ -4557,7 +4557,7 @@ query TT
 explain SELECT age FROM (SELECT * FROM person a join person b using (id, age, state));
 ----
 logical_plan
-01)Projection: a.age AS age
+01)Projection: a.age
 02)--Inner Join: a.id = b.id, a.age = b.age, a.state = b.state
 03)----SubqueryAlias: a
 04)------TableScan: person projection=[id, age, state]
@@ -4601,7 +4601,7 @@ query TT
 explain SELECT * FROM person a join person b using (id, age, state) join person c using (id, age, state);
 ----
 logical_plan
-01)Projection: a.id AS id, a.age AS age, a.state AS state
+01)Projection: a.id, a.age, a.state
 02)--Inner Join: a.id = c.id, a.age = c.age, a.state = c.state
 03)----Projection: a.id, a.age, a.state
 04)------Inner Join: a.id = b.id, a.age = b.age, a.state = b.state
@@ -4636,7 +4636,7 @@ query TT
 explain SELECT * FROM lineitem JOIN lineitem as lineitem2 USING (c1)
 ----
 logical_plan
-01)Projection: lineitem.c1 AS c1
+01)Projection: lineitem.c1
 02)--Inner Join: lineitem.c1 = lineitem2.c1
 03)----TableScan: lineitem projection=[c1]
 04)----SubqueryAlias: lineitem2