Skip to content

Commit 9c08afd

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 cee7bae commit 9c08afd

12 files changed

Lines changed: 596 additions & 33 deletions

File tree

‎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: 66 additions & 5 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;
2728
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;
@@ -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`].
6974
pub 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
84145
pub 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
}

‎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 {
@@ -857,12 +861,10 @@ impl LogicalPlanBuilder {
857861
// remove pushed down sort columns
858862
let new_expr = schema.columns().into_iter().map(Expr::Column).collect();
859863

864+
let input = Arc::unwrap_or_clone(self.plan);
865+
860866
let is_distinct = false;
861-
let plan = Self::add_missing_columns(
862-
Arc::unwrap_or_clone(self.plan),
863-
&missing_cols,
864-
is_distinct,
865-
)?;
867+
let plan = Self::add_missing_columns(input, &missing_cols, is_distinct)?;
866868

867869
let sort_plan = LogicalPlan::Sort(Sort {
868870
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
@@ -544,6 +544,35 @@ impl LogicalPlan {
544544
Ok(using_columns)
545545
}
546546

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

‎datafusion/expr/src/utils.rs‎

Lines changed: 49 additions & 1 deletion
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
};
@@ -440,6 +442,38 @@ fn exclude_using_columns(plan: &LogicalPlan) -> Result<HashSet<Column>> {
440442
Ok(excluded)
441443
}
442444

445+
fn merged_using_key_for_wildcard(
446+
col: Column,
447+
merged_keys: &[(Column, Column, JoinType)],
448+
) -> Result<Expr> {
449+
let Some((l, r, join_type)) =
450+
merged_keys.iter().find(|(l, r, _)| l == &col || r == &col)
451+
else {
452+
return Ok(Expr::Column(col));
453+
};
454+
455+
let (expr, visible_col) = match join_type {
456+
JoinType::Right | JoinType::RightSemi | JoinType::RightAnti => {
457+
(Expr::Column(r.clone()), r.clone())
458+
}
459+
JoinType::Full => CaseBuilder::new(
460+
None,
461+
vec![Expr::Column(l.clone()).is_not_null()],
462+
vec![Expr::Column(l.clone())],
463+
Some(Box::new(Expr::Column(r.clone()))),
464+
)
465+
.end()
466+
.map(|expr| (expr, col.clone()))?,
467+
_ => (Expr::Column(l.clone()), l.clone()),
468+
};
469+
470+
if expr == Expr::Column(visible_col.clone()) {
471+
Ok(expr)
472+
} else {
473+
Ok(expr.alias_qualified(visible_col.relation, visible_col.name))
474+
}
475+
}
476+
443477
/// Resolves an `Expr::Wildcard` to a collection of `Expr::Column`'s.
444478
pub fn expand_wildcard(
445479
schema: &DFSchema,
@@ -459,7 +493,21 @@ pub fn expand_wildcard(
459493
};
460494
// Add each excluded `Column` to columns_to_skip
461495
columns_to_skip.extend(excluded_columns);
462-
Ok(get_exprs_except_skipped(schema, &columns_to_skip))
496+
let exprs = get_exprs_except_skipped(schema, &columns_to_skip);
497+
498+
// Resolve the surviving USING / NATURAL key column to the merged key value,
499+
// while preserving the wildcard-visible field qualifier.
500+
let merged_keys = plan.using_key_pairs()?;
501+
if merged_keys.is_empty() {
502+
return Ok(exprs);
503+
}
504+
exprs
505+
.into_iter()
506+
.map(|expr| match expr {
507+
Expr::Column(col) => merged_using_key_for_wildcard(col, &merged_keys),
508+
other => Ok(other),
509+
})
510+
.collect()
463511
}
464512

465513
/// 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)