@@ -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 ;
27- use crate :: { Expr , ExprSchemable , LogicalPlan } ;
28+ use 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 ;
@@ -65,6 +66,10 @@ pub trait FunctionRewrite: Debug {
6566}
6667
6768/// Recursively normalize all [`Column`] expressions in the `expr` expression tree.
69+ ///
70+ /// An unqualified reference to the merged key of a `USING` / `NATURAL` join is
71+ /// resolved to a reserved internal merged-key expression via
72+ /// [`merged_using_key_or_column`].
6873pub fn normalize_col ( expr : Expr , plan : & LogicalPlan ) -> Result < Expr > {
6974 ColumnNormalizer :: new ( plan) . normalize ( expr)
7075}
@@ -74,13 +79,15 @@ pub fn normalize_col(expr: Expr, plan: &LogicalPlan) -> Result<Expr> {
7479pub ( crate ) struct ColumnNormalizer < ' a > {
7580 plan : & ' a LogicalPlan ,
7681 context : Option < ( Vec < & ' a DFSchema > , Vec < HashSet < Column > > ) > ,
82+ merged_keys : Option < Vec < ( Column , Column , JoinType ) > > ,
7783}
7884
7985impl < ' a > ColumnNormalizer < ' a > {
8086 pub ( crate ) fn new ( plan : & ' a LogicalPlan ) -> Self {
8187 Self {
8288 plan,
8389 context : None ,
90+ merged_keys : None ,
8491 }
8592 }
8693
@@ -105,15 +112,76 @@ impl<'a> ColumnNormalizer<'a> {
105112
106113 pub ( crate ) fn normalize ( & mut self , expr : Expr ) -> Result < Expr > {
107114 expr. transform ( |expr| match expr {
108- Expr :: Column ( column) => self
109- . normalize_column ( column)
110- . map ( |column| Transformed :: yes ( Expr :: Column ( column) ) ) ,
115+ Expr :: Column ( column) => {
116+ let was_unqualified = column. relation . is_none ( ) ;
117+ let column = self . normalize_column ( column) ?;
118+ let expr = if was_unqualified {
119+ let merged_keys = match & mut self . merged_keys {
120+ Some ( keys) => keys,
121+ keys @ None => keys. insert ( self . plan . using_key_pairs ( ) ?) ,
122+ } ;
123+ merged_using_key_or_column ( column, true , merged_keys) ?
124+ } else {
125+ Expr :: Column ( column)
126+ } ;
127+ Ok ( Transformed :: yes ( expr) )
128+ }
111129 _ => Ok ( Transformed :: no ( expr) ) ,
112130 } )
113131 . data ( )
114132 }
115133}
116134
135+ fn alias_merged_key ( expr : Expr , name : & str ) -> Expr {
136+ expr. alias_qualified (
137+ Some ( TableReference :: bare (
138+ LogicalPlanBuilder :: MERGED_KEY_QUALIFIER ,
139+ ) ) ,
140+ name,
141+ )
142+ }
143+
144+ /// Resolve a normalized column to itself, or -- when it is the merged key of a
145+ /// `USING` / `NATURAL` join referenced *unqualified* -- to the merged key
146+ /// expression with a reserved internal qualifier.
147+ ///
148+ /// `merged_keys` carries the `(left, right, join_type)` triples of the plan's
149+ /// USING / NATURAL joins. The merged key resolves to the left key for
150+ /// left-preserving joins, the right key for right-preserving joins, and
151+ /// `COALESCE(left, right)` for `FULL`. The COALESCE is built as
152+ /// `CASE WHEN left IS NOT NULL THEN left ELSE right END` (the form `coalesce`
153+ /// is simplified to, buildable here without depending on the functions crate).
154+ pub fn merged_using_key_or_column (
155+ col : Column ,
156+ was_unqualified : bool ,
157+ merged_keys : & [ ( Column , Column , JoinType ) ] ,
158+ ) -> Result < Expr > {
159+ if was_unqualified
160+ && let Some ( ( l, r, join_type) ) =
161+ merged_keys. iter ( ) . find ( |( l, r, _) | l == & col || r == & col)
162+ {
163+ let name = col. name . clone ( ) ;
164+ let expr = match join_type {
165+ // The left key is NULL-padded on right-only rows; the right key is
166+ // always present, so the merged key is just the right column.
167+ JoinType :: Right | JoinType :: RightSemi | JoinType :: RightAnti => {
168+ Expr :: Column ( r. clone ( ) )
169+ }
170+ // Either side may be NULL-padded; coalesce to the present one.
171+ JoinType :: Full => CaseBuilder :: new (
172+ None ,
173+ vec ! [ Expr :: Column ( l. clone( ) ) . is_not_null( ) ] ,
174+ vec ! [ Expr :: Column ( l. clone( ) ) ] ,
175+ Some ( Box :: new ( Expr :: Column ( r. clone ( ) ) ) ) ,
176+ )
177+ . end ( ) ?,
178+ _ => Expr :: Column ( l. clone ( ) ) ,
179+ } ;
180+ return Ok ( alias_merged_key ( expr, & name) ) ;
181+ }
182+ Ok ( Expr :: Column ( col) )
183+ }
184+
117185/// See [`Column::normalize_with_schemas_and_ambiguity_check`] for usage
118186pub fn normalize_col_with_schemas_and_ambiguity_check (
119187 expr : Expr ,
@@ -136,9 +204,9 @@ pub fn normalize_col_with_schemas_and_ambiguity_check(
136204 expr. transform ( |expr| {
137205 Ok ( {
138206 if let Expr :: Column ( c) = expr {
139- let col =
140- c. normalize_with_schemas_and_ambiguity_check ( schemas, using_columns) ?;
141- Transformed :: yes ( Expr :: Column ( col ) )
207+ Transformed :: yes ( Expr :: Column (
208+ c. normalize_with_schemas_and_ambiguity_check ( schemas, using_columns) ?,
209+ ) )
142210 } else {
143211 Transformed :: no ( expr)
144212 }
@@ -561,18 +629,24 @@ mod test {
561629 let right =
562630 crate :: logical_plan:: table_scan ( Some ( "right" ) , & schema, None ) ?. build ( ) ?;
563631 let plan = crate :: logical_plan:: table_scan ( Some ( "left" ) , & schema, None ) ?
564- . join_using ( right, crate :: JoinType :: Inner , vec ! [ Column :: from_name( "a" ) ] ) ?
632+ . join_using ( right, JoinType :: Inner , vec ! [ Column :: from_name( "a" ) ] ) ?
565633 . build ( ) ?;
634+ let merged_key = col ( "left.a" ) . alias_qualified (
635+ Some ( TableReference :: bare (
636+ LogicalPlanBuilder :: MERGED_KEY_QUALIFIER ,
637+ ) ) ,
638+ "a" ,
639+ ) ;
566640 assert_eq ! (
567641 super :: normalize_cols( [ col( "a" ) , col( "a" ) + col( "right.a" ) ] , & plan) ?,
568- vec![ col ( "left.a" ) , col ( "left.a" ) + col( "right.a" ) ]
642+ vec![ merged_key . clone ( ) , merged_key . clone ( ) + col( "right.a" ) ]
569643 ) ;
570644 let projected = LogicalPlanBuilder :: from ( plan. clone ( ) )
571645 . project ( [ col ( "a" ) . alias ( "key" ) , col ( "left.b" ) . alias ( "value" ) ] ) ?
572646 . build ( ) ?;
573647 assert_eq ! (
574648 projected. expressions( ) ,
575- vec![ col ( "left.a" ) . alias( "key" ) , col( "left.b" ) . alias( "value" ) ]
649+ vec![ merged_key . alias( "key" ) , col( "left.b" ) . alias( "value" ) ]
576650 ) ;
577651 let err = super :: normalize_cols ( [ col ( "a" ) , col ( "b" ) , col ( "missing" ) ] , & plan)
578652 . unwrap_err ( ) ;
@@ -589,7 +663,7 @@ mod test {
589663 let right =
590664 crate :: logical_plan:: table_scan ( Some ( "right" ) , & schema, None ) ?. build ( ) ?;
591665 let mut plan = crate :: logical_plan:: table_scan ( Some ( "left" ) , & schema, None ) ?
592- . join_using ( right, crate :: JoinType :: Inner , vec ! [ Column :: from_name( "a" ) ] ) ?
666+ . join_using ( right, JoinType :: Inner , vec ! [ Column :: from_name( "a" ) ] ) ?
593667 . build ( ) ?;
594668 // Invalid USING keys must only be inspected when an unqualified column
595669 // needs normalization, just as with a single expression.
0 commit comments