@@ -22,14 +22,15 @@ use std::collections::HashSet;
2222use std:: fmt:: Debug ;
2323use std:: sync:: Arc ;
2424
25+ use crate :: conditional_expressions:: CaseBuilder ;
2526use crate :: expr:: { Alias , Sort , Unnest } ;
2627use crate :: logical_plan:: Projection ;
2728use crate :: { Expr , ExprSchemable , LogicalPlan , LogicalPlanBuilder } ;
2829
2930use datafusion_common:: TableReference ;
3031use datafusion_common:: config:: ConfigOptions ;
3132use datafusion_common:: tree_node:: { Transformed , TransformedResult , TreeNode } ;
32- use datafusion_common:: { Column , DFSchema , Result } ;
33+ use datafusion_common:: { Column , DFSchema , JoinType , Result } ;
3334
3435mod guarantees;
3536pub use guarantees:: GuaranteeRewriter ;
@@ -66,12 +67,22 @@ pub trait FunctionRewrite: Debug {
6667
6768/// Recursively call `LogicalPlanBuilder::normalize` on all [`Column`] expressions
6869/// in the `expr` expression tree.
70+ ///
71+ /// An unqualified reference to the merged key of a `USING` / `NATURAL` join is
72+ /// resolved to a reserved internal merged-key expression via
73+ /// [`merged_using_key_or_column`].
6974pub fn normalize_col ( expr : Expr , plan : & LogicalPlan ) -> Result < Expr > {
75+ let merged_keys = plan. using_key_pairs ( ) ?;
7076 expr. transform ( |expr| {
7177 Ok ( {
7278 if let Expr :: Column ( c) = expr {
79+ let was_unqualified = c. relation . is_none ( ) ;
7380 let col = LogicalPlanBuilder :: normalize ( plan, c) ?;
74- Transformed :: yes ( Expr :: Column ( col) )
81+ Transformed :: yes ( merged_using_key_or_column (
82+ col,
83+ was_unqualified,
84+ & merged_keys,
85+ ) ?)
7586 } else {
7687 Transformed :: no ( expr)
7788 }
@@ -80,6 +91,56 @@ pub fn normalize_col(expr: Expr, plan: &LogicalPlan) -> Result<Expr> {
8091 . data ( )
8192}
8293
94+ fn alias_merged_key ( expr : Expr , name : & str ) -> Expr {
95+ expr. alias_qualified (
96+ Some ( TableReference :: bare (
97+ LogicalPlanBuilder :: MERGED_KEY_QUALIFIER ,
98+ ) ) ,
99+ name,
100+ )
101+ }
102+
103+ /// Resolve a normalized column to itself, or -- when it is the merged key of a
104+ /// `USING` / `NATURAL` join referenced *unqualified* -- to the merged key
105+ /// expression with a reserved internal qualifier.
106+ ///
107+ /// `merged_keys` carries the `(left, right, join_type)` triples of the plan's
108+ /// USING / NATURAL joins. The merged key resolves to the left key for
109+ /// left-preserving joins, the right key for right-preserving joins, and
110+ /// `COALESCE(left, right)` for `FULL`. The COALESCE is built as
111+ /// `CASE WHEN left IS NOT NULL THEN left ELSE right END` (the form `coalesce`
112+ /// is simplified to, buildable here without depending on the functions crate).
113+ pub fn merged_using_key_or_column (
114+ col : Column ,
115+ was_unqualified : bool ,
116+ merged_keys : & [ ( Column , Column , JoinType ) ] ,
117+ ) -> Result < Expr > {
118+ if was_unqualified
119+ && let Some ( ( l, r, join_type) ) =
120+ merged_keys. iter ( ) . find ( |( l, r, _) | l == & col || r == & col)
121+ {
122+ let name = col. name . clone ( ) ;
123+ let expr = match join_type {
124+ // The left key is NULL-padded on right-only rows; the right key is
125+ // always present, so the merged key is just the right column.
126+ JoinType :: Right | JoinType :: RightSemi | JoinType :: RightAnti => {
127+ Expr :: Column ( r. clone ( ) )
128+ }
129+ // Either side may be NULL-padded; coalesce to the present one.
130+ JoinType :: Full => CaseBuilder :: new (
131+ None ,
132+ vec ! [ Expr :: Column ( l. clone( ) ) . is_not_null( ) ] ,
133+ vec ! [ Expr :: Column ( l. clone( ) ) ] ,
134+ Some ( Box :: new ( Expr :: Column ( r. clone ( ) ) ) ) ,
135+ )
136+ . end ( ) ?,
137+ _ => Expr :: Column ( l. clone ( ) ) ,
138+ } ;
139+ return Ok ( alias_merged_key ( expr, & name) ) ;
140+ }
141+ Ok ( Expr :: Column ( col) )
142+ }
143+
83144/// See [`Column::normalize_with_schemas_and_ambiguity_check`] for usage
84145pub fn normalize_col_with_schemas_and_ambiguity_check (
85146 expr : Expr ,
@@ -102,9 +163,9 @@ pub fn normalize_col_with_schemas_and_ambiguity_check(
102163 expr. transform ( |expr| {
103164 Ok ( {
104165 if let Expr :: Column ( c) = expr {
105- let col =
106- c. normalize_with_schemas_and_ambiguity_check ( schemas, using_columns) ?;
107- Transformed :: yes ( Expr :: Column ( col ) )
166+ Transformed :: yes ( Expr :: Column (
167+ c. normalize_with_schemas_and_ambiguity_check ( schemas, using_columns) ?,
168+ ) )
108169 } else {
109170 Transformed :: no ( expr)
110171 }
0 commit comments