Skip to content

Commit b89ffe3

Browse files
committed
fix: order by a qualified key alongside a FULL USING/NATURAL merged key
A FULL USING / NATURAL join exposes its merged key as an unqualified COALESCE(left, right) column. Ordering by a qualified key of the same join (`ORDER BY a.k`) folds that qualified `a.k` into the projection that already carries the unqualified merged `k`, building an illegal `{ k, a.k }` schema, so such queries errored at plan time. During sort push-down, rename a colliding merged key to a reserved internal name before folding the qualified sort column, and restore the original name in the wrapper projection. References to the merged key in the sort itself (`ORDER BY a.k, k`) are rewritten to the same internal name so multi-key sorts keep resolving; that shape was legal before the merged-key work, so this avoids a regression. Teach the unparser to flatten a wrapper projection that purely renames an inner column over a Sort, so the rewritten plan still round-trips through plan -> SQL -> plan. Refs #22881 Signed-off-by: Jiawei Zhao <Phoenix500526@163.com>
1 parent 98518d2 commit b89ffe3

4 files changed

Lines changed: 263 additions & 8 deletions

File tree

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

Lines changed: 151 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ use arrow::datatypes::{DataType, Field, FieldRef, Fields, Schema, SchemaRef};
5353
use datafusion_common::display::ToStringifiedPlan;
5454
use datafusion_common::file_options::file_type::FileType;
5555
use datafusion_common::metadata::FieldMetadata;
56+
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
5657
use datafusion_common::{
5758
Column, Constraints, DFSchema, DFSchemaRef, NullEquality, Result, ScalarValue,
5859
TableReference, ToDFSchema, UnnestOptions, exec_err,
@@ -130,6 +131,11 @@ pub struct LogicalPlanBuilder {
130131
}
131132

132133
impl LogicalPlanBuilder {
134+
/// Reserved name prefix for a FULL USING / NATURAL join merged key that has
135+
/// been renamed to dodge an unqualified/qualified schema collision during
136+
/// sort push-down. A real column carrying this prefix would clash.
137+
const MERGED_KEY_NAME_PREFIX: &str = "__datafusion_merged_key_";
138+
133139
/// Create a builder from an existing plan
134140
pub fn new(plan: LogicalPlan) -> Self {
135141
Self {
@@ -819,6 +825,97 @@ impl LogicalPlanBuilder {
819825
self.sort_with_limit(sorts, None)
820826
}
821827

828+
/// Internal name for a FULL USING / NATURAL join merged key that has been
829+
/// renamed to dodge an unqualified/qualified schema collision during sort
830+
/// push-down. See [`Self::sort_with_limit`].
831+
fn synthetic_merged_name(name: &str) -> String {
832+
format!("{}{name}", Self::MERGED_KEY_NAME_PREFIX)
833+
}
834+
835+
/// Unqualified projection output fields that collide with a *qualified*
836+
/// missing sort column of the same name. In practice these are the
837+
/// `COALESCE` merged keys of FULL USING / NATURAL joins, exposed
838+
/// unqualified, which `ORDER BY <qualified key>` would otherwise drag into
839+
/// an illegal `{ k, a.k }` projection schema. They are renamed to a
840+
/// synthetic name before folding; any unqualified reference to them in the
841+
/// sort itself is rewritten to match (see [`Self::rewrite_sort_merged_keys`]).
842+
fn renamable_merged_keys(
843+
schema: &DFSchemaRef,
844+
missing_cols: &IndexSet<Column>,
845+
) -> HashSet<String> {
846+
let qualified_missing: HashSet<&str> = missing_cols
847+
.iter()
848+
.filter(|c| c.relation.is_some())
849+
.map(|c| c.name.as_str())
850+
.collect();
851+
if qualified_missing.is_empty() {
852+
return HashSet::new();
853+
}
854+
855+
schema
856+
.columns()
857+
.into_iter()
858+
.filter(|c| {
859+
c.relation.is_none() && qualified_missing.contains(c.name.as_str())
860+
})
861+
.map(|c| c.name)
862+
.collect()
863+
}
864+
865+
/// Re-alias a renamable merged key output field to its
866+
/// [`Self::synthetic_merged_name`] so it no longer collides with the
867+
/// qualified sort column folded alongside it.
868+
fn rename_merged_key(expr: Expr, renamable: &HashSet<String>) -> Expr {
869+
match expr {
870+
Expr::Alias(mut alias)
871+
if alias.relation.is_none() && renamable.contains(&alias.name) =>
872+
{
873+
alias.name = Self::synthetic_merged_name(&alias.name);
874+
Expr::Alias(alias)
875+
}
876+
Expr::Column(column)
877+
if column.relation.is_none() && renamable.contains(&column.name) =>
878+
{
879+
let name = Self::synthetic_merged_name(&column.name);
880+
Expr::Column(column).alias(name)
881+
}
882+
other => other,
883+
}
884+
}
885+
886+
/// Rewrite unqualified references to a renamed merged key inside the sort
887+
/// expressions to its [`Self::synthetic_merged_name`], so a sort that also
888+
/// orders by the merged key (e.g. `ORDER BY a.k, k`) keeps resolving after
889+
/// the projection field has been renamed.
890+
fn rewrite_sort_merged_keys(
891+
sorts: Vec<SortExpr>,
892+
renamable: &HashSet<String>,
893+
) -> Result<Vec<SortExpr>> {
894+
sorts
895+
.into_iter()
896+
.map(|sort| {
897+
let expr = sort
898+
.expr
899+
.clone()
900+
.transform(|e| {
901+
Ok(match e {
902+
Expr::Column(c)
903+
if c.relation.is_none()
904+
&& renamable.contains(&c.name) =>
905+
{
906+
Transformed::yes(Expr::Column(Column::new_unqualified(
907+
Self::synthetic_merged_name(&c.name),
908+
)))
909+
}
910+
other => Transformed::no(other),
911+
})
912+
})
913+
.data()?;
914+
Ok(sort.with_expr(expr))
915+
})
916+
.collect()
917+
}
918+
822919
/// Apply a sort
823920
pub fn sort_with_limit(
824921
self,
@@ -852,15 +949,63 @@ impl LogicalPlanBuilder {
852949
})));
853950
}
854951

952+
// A qualified sort column folded into the projection can collide with an
953+
// unqualified output field of the same name: the COALESCE merged key of a
954+
// FULL USING / NATURAL join is exposed unqualified, so `ORDER BY a.k`
955+
// dragging the qualified `a.k` into the `SELECT k` projection would build
956+
// an illegal `{ k, a.k }` schema. Rename such merged keys to a unique
957+
// internal name before folding, and restore the original name below in
958+
// the wrapper projection.
959+
let renamable = match self.plan.as_ref() {
960+
LogicalPlan::Projection(_) => {
961+
Self::renamable_merged_keys(schema, &missing_cols)
962+
}
963+
_ => HashSet::new(),
964+
};
965+
966+
// Keep the sort's own references to a renamed merged key resolving by
967+
// rewriting them to the same synthetic name (e.g. `ORDER BY a.k, k`,
968+
// where `k` is the merged key being renamed to dodge the `a.k` clash).
969+
let sorts = if renamable.is_empty() {
970+
sorts
971+
} else {
972+
Self::rewrite_sort_merged_keys(sorts, &renamable)?
973+
};
974+
855975
// remove pushed down sort columns
856-
let new_expr = schema.columns().into_iter().map(Expr::Column).collect();
976+
let new_expr = schema
977+
.columns()
978+
.into_iter()
979+
.map(|c| {
980+
if c.relation.is_none() && renamable.contains(&c.name) {
981+
Expr::Column(Column::new_unqualified(Self::synthetic_merged_name(
982+
&c.name,
983+
)))
984+
.alias(c.name)
985+
} else {
986+
Expr::Column(c)
987+
}
988+
})
989+
.collect();
990+
991+
let input = if renamable.is_empty() {
992+
Arc::unwrap_or_clone(self.plan)
993+
} else {
994+
match Arc::unwrap_or_clone(self.plan) {
995+
LogicalPlan::Projection(Projection { expr, input, .. }) => {
996+
let expr = expr
997+
.into_iter()
998+
.map(|e| Self::rename_merged_key(e, &renamable))
999+
.collect::<Vec<_>>();
1000+
LogicalPlan::Projection(Projection::try_new(expr, input)?)
1001+
}
1002+
// `renamable` is only non-empty for a Projection (checked above).
1003+
other => other,
1004+
}
1005+
};
8571006

8581007
let is_distinct = false;
859-
let plan = Self::add_missing_columns(
860-
Arc::unwrap_or_clone(self.plan),
861-
&missing_cols,
862-
is_distinct,
863-
)?;
1008+
let plan = Self::add_missing_columns(input, &missing_cols, is_distinct)?;
8641009

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

‎datafusion/sql/src/unparser/rewrite.rs‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,20 @@ pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
220220
})
221221
.collect::<Vec<_>>();
222222

223-
let mut collects = p.expr.clone();
223+
let mut collects = p
224+
.expr
225+
.iter()
226+
.map(|e| match e {
227+
// A pure rename `Column(x) AS out` compares as its underlying
228+
// column `x`, so an outer Projection that only renames an inner
229+
// column (e.g. the restored merged key of a FULL USING / NATURAL
230+
// join) still matches and can be flattened instead of nested.
231+
Expr::Alias(alias) if matches!(alias.expr.as_ref(), Expr::Column(_)) => {
232+
alias.expr.as_ref().clone()
233+
}
234+
_ => e.clone(),
235+
})
236+
.collect::<Vec<_>>();
224237
for sort in &sort.expr {
225238
// Strip aliases from sort expressions so the comparison matches
226239
// the inner Projection's raw expressions. The optimizer may add
@@ -250,7 +263,19 @@ pub(super) fn rewrite_plan_for_sort_on_non_projected_fields(
250263
let new_exprs = p
251264
.expr
252265
.iter()
253-
.map(|e| map.get(e).unwrap_or(e).clone())
266+
.map(|e| {
267+
// For a pure rename `Column(x) AS out`, inline the inner
268+
// definition of `x` and re-apply the outer output name, so the
269+
// flattened SELECT keeps `out` as the column name.
270+
if let Expr::Alias(alias) = e
271+
&& matches!(alias.expr.as_ref(), Expr::Column(_))
272+
&& let Some(inner) = map.get(alias.expr.as_ref())
273+
{
274+
inner.clone().unalias().alias(alias.name.clone())
275+
} else {
276+
map.get(e).unwrap_or(e).clone()
277+
}
278+
})
254279
.collect::<Vec<_>>();
255280

256281
// The inner Projection may define aliases that the Sort references

‎datafusion/sql/tests/cases/plan_to_sql.rs‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ use datafusion_sql::unparser::dialect::{
4646
};
4747
use datafusion_sql::unparser::{Unparser, expr_to_sql, plan_to_sql};
4848
use insta::assert_snapshot;
49+
use rstest::rstest;
4950
use sqlparser::ast::Statement;
5051
use std::hash::Hash;
5152
use std::ops::Add;
@@ -260,6 +261,48 @@ fn roundtrip_statement() -> Result<()> {
260261
Ok(())
261262
}
262263

264+
/// Regression for https://github.com/apache/datafusion/issues/22881.
265+
///
266+
/// A FULL JOIN USING / NATURAL merged key is `COALESCE(left, right)`, exposed
267+
/// as an *unqualified* column. Selecting that merged key while ordering by a
268+
/// *qualified* key of the same join (`ORDER BY ta.id`) must both plan and
269+
/// survive an unparser round-trip (`plan -> SQL -> plan`). Earlier fixes either
270+
/// errored at plan time or produced SQL that could not be re-planned. `USING`
271+
/// and `NATURAL` are covered separately since the unparser renders them
272+
/// differently.
273+
#[rstest]
274+
#[case::full_join_using(
275+
"SELECT id FROM (SELECT j1_id AS id FROM j1) ta \
276+
FULL JOIN (SELECT j2_id AS id FROM j2) tb USING (id) \
277+
ORDER BY ta.id"
278+
)]
279+
#[case::natural_full_join(
280+
"SELECT id FROM (SELECT j1_id AS id FROM j1) ta \
281+
NATURAL FULL JOIN (SELECT j2_id AS id FROM j2) tb \
282+
ORDER BY ta.id"
283+
)]
284+
fn roundtrip_full_join_merged_key_order_by_qualified(#[case] query: &str) -> Result<()> {
285+
let dialect = GenericDialect {};
286+
let statement = Parser::new(&dialect)
287+
.try_with_sql(query)?
288+
.parse_statement()?;
289+
290+
let context = MockContextProvider {
291+
state: MockSessionState::default(),
292+
};
293+
let sql_to_rel = SqlToRel::new(&context);
294+
let plan = sql_to_rel.sql_statement_to_plan(statement).unwrap();
295+
296+
let roundtrip_statement = plan_to_sql(&plan)?;
297+
let plan_roundtrip = sql_to_rel
298+
.sql_statement_to_plan(roundtrip_statement)
299+
.unwrap();
300+
301+
assert_eq!(plan, plan_roundtrip);
302+
303+
Ok(())
304+
}
305+
263306
#[test]
264307
fn roundtrip_crossjoin() -> Result<()> {
265308
let query = "select j1.j1_id, j2.j2_string from j1, j2";

‎datafusion/sqllogictest/test_files/join_using_merged_key.slt‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,48 @@ SELECT coalesce(a.k, b.k) AS k FROM a FULL JOIN b ON a.k = b.k ORDER BY k
119119
3
120120
4
121121

122+
##########
123+
# Regression: ORDER BY a *qualified* key while selecting the FULL merged key.
124+
# The merged k is COALESCE(a.k, b.k), exposed unqualified; ORDER BY a.k folds
125+
# the qualified a.k into the same projection. This must plan (it regressed to a
126+
# schema-ambiguity error) and sort by a.k, not by the merged key: the right-only
127+
# row (a.k IS NULL) sorts last, the rest descend by a.k -- so the output differs
128+
# from ORDER BY k DESC (which would be 4, 3, 2, 1).
129+
##########
130+
131+
query I nosort
132+
SELECT k FROM a FULL JOIN b USING (k) ORDER BY a.k DESC NULLS LAST
133+
----
134+
3
135+
2
136+
1
137+
4
138+
139+
query I nosort
140+
SELECT k FROM a NATURAL FULL JOIN b ORDER BY a.k DESC NULLS LAST
141+
----
142+
3
143+
2
144+
1
145+
4
146+
147+
##########
148+
# Regression: a *qualified* key and the *unqualified* merged key together in one
149+
# ORDER BY (`ORDER BY a.k DESC NULLS LAST, k`). On main this was legal (the merged
150+
# key resolved to a.k, so both sort keys were a.k); the merged-key fix must keep
151+
# it legal -- renaming the merged key to dodge the `{k, a.k}` collision must also
152+
# rewrite the unqualified `k` reference in the sort. a.k is unique here so the
153+
# secondary `k` never breaks a tie, but the shape still exercises that path.
154+
##########
155+
156+
query I nosort
157+
SELECT k FROM a FULL JOIN b USING (k) ORDER BY a.k DESC NULLS LAST, k
158+
----
159+
3
160+
2
161+
1
162+
4
163+
122164
##########
123165
# Wildcard over the join: the merged key must appear once and carry the right
124166
# value. (commit 2 target -- the materialized merged column must dedupe under *)

0 commit comments

Comments
 (0)