Skip to content

Commit 7059068

Browse files
committed
fix(sql): resolve USING/NATURAL merged keys
Unqualified USING and NATURAL keys could return NULL for right-only rows or fail beside qualified side keys. Resolve merged keys by join type across SELECT, WHERE, ORDER BY and wildcards while preserving qualified side access and SQL round trips. CLOSES #22881 Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
1 parent d137137 commit 7059068

14 files changed

Lines changed: 637 additions & 71 deletions

File tree

‎datafusion/core/src/datasource/view_test.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -289,10 +289,10 @@ mod tests {
289289

290290
insta::assert_snapshot!(batches_to_string(&results),@r"
291291
+---------+---------+---------+
292-
| column2 | column1 | column3 |
292+
| column1 | column2 | column3 |
293293
+---------+---------+---------+
294-
| 2 | 1 | 3 |
295-
| 5 | 4 | 6 |
294+
| 1 | 2 | 3 |
295+
| 4 | 5 | 6 |
296296
+---------+---------+---------+
297297
");
298298

‎datafusion/core/tests/sql/unparser.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -378,7 +378,7 @@ async fn optimized_duckdb_unparse_qualifies_nested_passthrough_column() -> Resul
378378
// `o` (which is only the base-table alias one level deeper). The bug emitted
379379
// `"o"."order_id"` inside that derived table; the fix emits a bare column.
380380
let expected = concat!(
381-
r#"SELECT "o"."order_id", "o"."discount_pct_2" "#,
381+
r#"SELECT "oi"."order_id", "o"."discount_pct_2" "#,
382382
r#"FROM "warehouse"."main"."order_items" AS "oi" "#,
383383
r#"INNER JOIN (SELECT "order_id", "#,
384384
r#"CASE WHEN "__common_expr_1" IS NOT NULL "#,

‎datafusion/expr/src/expr_rewriter/mod.rs‎

Lines changed: 86 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,15 @@ use std::collections::HashSet;
2222
use std::fmt::Debug;
2323
use std::sync::Arc;
2424

25+
use crate::conditional_expressions::CaseBuilder;
2526
use crate::expr::{Alias, Sort, Unnest};
2627
use crate::logical_plan::Projection;
27-
use crate::{Expr, ExprSchemable, LogicalPlan};
28+
use crate::{Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder};
2829

2930
use datafusion_common::TableReference;
3031
use datafusion_common::config::ConfigOptions;
3132
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
32-
use datafusion_common::{Column, DFSchema, Result};
33+
use datafusion_common::{Column, DFSchema, JoinType, Result};
3334

3435
mod guarantees;
3536
pub 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`].
6873
pub 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> {
7479
pub(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

7985
impl<'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
118186
pub 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.

‎datafusion/expr/src/logical_plan/builder.rs‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,10 @@ pub struct LogicalPlanBuilder {
132132
}
133133

134134
impl LogicalPlanBuilder {
135+
/// Reserved internal qualifier for USING / NATURAL join merged keys.
136+
/// A real relation carrying this qualifier would clash.
137+
pub const MERGED_KEY_QUALIFIER: &str = "__datafusion_merged_key_qualifier__";
138+
135139
/// Create a builder from an existing plan
136140
pub fn new(plan: LogicalPlan) -> Self {
137141
Self {
@@ -854,12 +858,10 @@ impl LogicalPlanBuilder {
854858
// remove pushed down sort columns
855859
let new_expr = schema.columns().into_iter().map(Expr::Column).collect();
856860

861+
let input = Arc::unwrap_or_clone(self.plan);
862+
857863
let is_distinct = false;
858-
let plan = Self::add_missing_columns(
859-
Arc::unwrap_or_clone(self.plan),
860-
&missing_cols,
861-
is_distinct,
862-
)?;
864+
let plan = Self::add_missing_columns(input, &missing_cols, is_distinct)?;
863865

864866
let sort_plan = LogicalPlan::Sort(Sort {
865867
expr: normalize_sorts(sorts, &plan)?,

‎datafusion/expr/src/logical_plan/plan.rs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -546,6 +546,35 @@ impl LogicalPlan {
546546
Ok(using_columns)
547547
}
548548

549+
/// Returns the `(left, right, join_type)` join-key column triples of every
550+
/// `USING` / `NATURAL` join in this plan.
551+
pub fn using_key_pairs(
552+
&self,
553+
) -> Result<Vec<(Column, Column, JoinType)>, DataFusionError> {
554+
let mut triples = vec![];
555+
556+
self.apply_with_subqueries(|plan| {
557+
if let LogicalPlan::Join(Join {
558+
join_constraint: JoinConstraint::Using,
559+
join_type,
560+
on,
561+
..
562+
}) = plan
563+
{
564+
for (l, r) in on {
565+
if let (Some(l), Some(r)) =
566+
(l.get_as_join_column(), r.get_as_join_column())
567+
{
568+
triples.push((l.to_owned(), r.to_owned(), *join_type));
569+
}
570+
}
571+
}
572+
Ok(TreeNodeRecursion::Continue)
573+
})?;
574+
575+
Ok(triples)
576+
}
577+
549578
/// returns the first output expression of this `LogicalPlan` node.
550579
pub fn head_output_expr(&self) -> Result<Option<Expr>> {
551580
match self {

‎datafusion/expr/src/utils.rs‎

Lines changed: 55 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,10 @@ use std::cmp::Ordering;
2121
use std::collections::{BTreeSet, HashSet};
2222
use std::sync::Arc;
2323

24+
use crate::conditional_expressions::CaseBuilder;
2425
use crate::expr::{Alias, Sort, WildcardOptions, WindowFunctionParams};
2526
use crate::expr_rewriter::strip_outer_reference;
27+
use crate::logical_plan::JoinType;
2628
use crate::{
2729
BinaryExpr, Expr, ExprSchemable, Filter, GroupingSet, LogicalPlan, Operator, and,
2830
};
@@ -409,37 +411,54 @@ fn get_exprs_except_skipped(
409411
/// joins). An unqualified wildcard should include each USING column only once.
410412
/// This function returns the duplicate columns that should be excluded.
411413
fn exclude_using_columns(plan: &LogicalPlan) -> Result<HashSet<Column>> {
412-
let output_columns: HashSet<_> = plan.schema().columns().iter().cloned().collect();
414+
let output_columns = plan.schema().columns();
413415
let mut excluded = HashSet::new();
414416
for cols in plan.using_columns()? {
415417
// `using_columns()` returns join columns from both sides regardless of
416-
// the join type. For semi/anti joins, only one side's columns appear in
417-
// the output schema. Filter to output columns so that columns from the
418-
// non-output side don't participate in the deduplication process below
419-
// and displace real output columns.
420-
let mut cols: Vec<_> = cols
421-
.into_iter()
422-
.filter(|c| output_columns.contains(c))
423-
.collect();
424-
425-
// Sort so we keep the same qualified column, regardless of HashSet
426-
// iteration order.
427-
cols.sort();
428-
429-
// Keep only one column per name from the columns set, adding any
430-
// duplicates to the excluded set.
418+
// the join type. Visit only output columns, in schema order, so the
419+
// surviving key keeps its position in an unqualified wildcard.
431420
let mut seen_names = HashSet::new();
432-
for col in cols {
433-
if seen_names.contains(col.name.as_str()) {
434-
excluded.insert(col); // exclude columns with already seen name
435-
} else {
436-
seen_names.insert(col.name.clone()); // mark column name as seen
421+
for col in output_columns.iter().filter(|col| cols.contains(*col)) {
422+
if !seen_names.insert(&col.name) {
423+
excluded.insert(col.clone());
437424
}
438425
}
439426
}
440427
Ok(excluded)
441428
}
442429

430+
fn merged_using_key_for_wildcard(
431+
col: Column,
432+
merged_keys: &[(Column, Column, JoinType)],
433+
) -> Result<Expr> {
434+
let Some((l, r, join_type)) =
435+
merged_keys.iter().find(|(l, r, _)| l == &col || r == &col)
436+
else {
437+
return Ok(Expr::Column(col));
438+
};
439+
440+
let (expr, visible_col) = match join_type {
441+
JoinType::Right | JoinType::RightSemi | JoinType::RightAnti => {
442+
(Expr::Column(r.clone()), r.clone())
443+
}
444+
JoinType::Full => CaseBuilder::new(
445+
None,
446+
vec![Expr::Column(l.clone()).is_not_null()],
447+
vec![Expr::Column(l.clone())],
448+
Some(Box::new(Expr::Column(r.clone()))),
449+
)
450+
.end()
451+
.map(|expr| (expr, col.clone()))?,
452+
_ => (Expr::Column(l.clone()), l.clone()),
453+
};
454+
455+
if expr == Expr::Column(visible_col.clone()) {
456+
Ok(expr)
457+
} else {
458+
Ok(expr.alias_qualified(visible_col.relation, visible_col.name))
459+
}
460+
}
461+
443462
/// Resolves an `Expr::Wildcard` to a collection of `Expr::Column`'s.
444463
pub fn expand_wildcard(
445464
schema: &DFSchema,
@@ -459,7 +478,21 @@ pub fn expand_wildcard(
459478
};
460479
// Add each excluded `Column` to columns_to_skip
461480
columns_to_skip.extend(excluded_columns);
462-
Ok(get_exprs_except_skipped(schema, &columns_to_skip))
481+
let exprs = get_exprs_except_skipped(schema, &columns_to_skip);
482+
483+
// Resolve the surviving USING / NATURAL key column to the merged key value,
484+
// while preserving the wildcard-visible field qualifier.
485+
let merged_keys = plan.using_key_pairs()?;
486+
if merged_keys.is_empty() {
487+
return Ok(exprs);
488+
}
489+
exprs
490+
.into_iter()
491+
.map(|expr| match expr {
492+
Expr::Column(col) => merged_using_key_for_wildcard(col, &merged_keys),
493+
other => Ok(other),
494+
})
495+
.collect()
463496
}
464497

465498
/// Resolves an `Expr::Wildcard` to a collection of qualified `Expr::Column`'s.

‎datafusion/optimizer/src/optimize_projections/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2035,7 +2035,7 @@ mod tests {
20352035
assert_snapshot!(
20362036
optimized_plan.clone(),
20372037
@r"
2038-
Projection: test.a, test.b
2038+
Projection: test.a AS a, test.b
20392039
Left Join: Using test.a = test2.a
20402040
TableScan: test projection=[a, b]
20412041
TableScan: test2 projection=[a]

0 commit comments

Comments
 (0)