Skip to content

Commit a03258c

Browse files
1fanwang1fanwang
committed
fix: keep input qualifiers in INTERSECT ALL and EXCEPT ALL
Both sides used the same name for their synthetic row number column, so the sides were always requalified as `left` and `right`, and the final projection rebuilt every output column as `left.<name>`. Queries that referred to the original qualifier after the set operation failed, and two columns with the same name from different relations on one side collapsed into duplicate `left.x` outputs. Give each side its own row number name so requalification only happens when the user's columns conflict, join on the columns the requalified schemas actually contain, and alias the output back to the left input's qualifiers and names. Signed-off-by: Stefan Wang <[email protected]>
1 parent 09c4eaa commit a03258c

3 files changed

Lines changed: 119 additions & 48 deletions

File tree

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

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,69 @@ async fn set_operation_all_preserves_synthetic_column_name() -> Result<()> {
103103
}
104104
Ok(())
105105
}
106+
107+
#[tokio::test]
108+
async fn set_operation_all_keeps_left_qualifiers() -> Result<()> {
109+
let ctx = SessionContext::new();
110+
ctx.sql("CREATE TABLE t (x INT) AS VALUES (1), (1), (2)")
111+
.await?
112+
.collect()
113+
.await?;
114+
ctx.sql("CREATE TABLE u (x INT) AS VALUES (1), (3)")
115+
.await?
116+
.collect()
117+
.await?;
118+
119+
let result = ctx
120+
.sql("SELECT t.x FROM t INTERSECT ALL SELECT u.x FROM u ORDER BY t.x")
121+
.await?
122+
.collect()
123+
.await?;
124+
assert_batches_sorted_eq!(["+---+", "| x |", "+---+", "| 1 |", "+---+"], &result);
125+
126+
let t = ctx.table("t").await?;
127+
let u = ctx.table("u").await?;
128+
let result = t
129+
.clone()
130+
.intersect(u.clone())?
131+
.select(vec![col("t.x")])?
132+
.collect()
133+
.await?;
134+
assert_batches_sorted_eq!(["+---+", "| x |", "+---+", "| 1 |", "+---+"], &result);
135+
136+
let result = t.except(u)?.select(vec![col("t.x")])?.collect().await?;
137+
assert_batches_sorted_eq!(
138+
["+---+", "| x |", "+---+", "| 1 |", "| 2 |", "+---+"],
139+
&result
140+
);
141+
Ok(())
142+
}
143+
144+
#[tokio::test]
145+
async fn set_operation_all_with_same_name_from_two_relations() -> Result<()> {
146+
let ctx = SessionContext::new();
147+
let side = "SELECT t1.x, t2.x FROM (VALUES (1)) t1(x) CROSS JOIN (VALUES (2)) t2(x)";
148+
let result = ctx
149+
.sql(&format!("{side} INTERSECT ALL {side}"))
150+
.await?
151+
.collect()
152+
.await?;
153+
assert_batches_sorted_eq!(
154+
[
155+
"+---+---+",
156+
"| x | x |",
157+
"+---+---+",
158+
"| 1 | 2 |",
159+
"+---+---+"
160+
],
161+
&result
162+
);
163+
164+
let result = ctx
165+
.sql(&format!("{side} EXCEPT ALL {side}"))
166+
.await?
167+
.collect()
168+
.await?;
169+
assert!(result.iter().all(|batch| batch.num_rows() == 0));
170+
Ok(())
171+
}

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

Lines changed: 47 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1549,77 +1549,78 @@ impl LogicalPlanBuilder {
15491549
&& left.data_type() == right.data_type()
15501550
})
15511551
.collect::<Vec<_>>();
1552+
// Each side gets its own row number column name, so the two synthetic
1553+
// columns never collide with each other. The sides are then only
1554+
// requalified when the user's own columns conflict.
15521555
let mut row_number_name = "__datafusion_set_operation_row_number".to_string();
1553-
while [left_plan.schema(), right_plan.schema()]
1554-
.iter()
1555-
.any(|schema| {
1556-
schema
1557-
.fields()
1558-
.iter()
1559-
.any(|field| field.name() == &row_number_name)
1560-
})
1556+
let name_in_use = |name: &str| {
1557+
[left_plan.schema(), right_plan.schema()]
1558+
.iter()
1559+
.any(|schema| schema.fields().iter().any(|field| field.name() == name))
1560+
};
1561+
while name_in_use(&row_number_name)
1562+
|| name_in_use(&format!("{row_number_name}_right"))
15611563
{
15621564
row_number_name.push('_');
15631565
}
1564-
let with_row_number = |plan: LogicalPlan, columns: &[Column]| {
1566+
let right_row_number_name = format!("{row_number_name}_right");
1567+
let with_row_number = |plan: LogicalPlan, columns: &[Column], name: &str| {
15651568
let mut window = crate::expr::WindowFunction::new(
15661569
crate::WindowFunctionDefinition::WindowUDF(Arc::clone(row_number)),
15671570
vec![],
15681571
);
15691572
window.params.partition_by =
15701573
columns.iter().cloned().map(Expr::Column).collect();
15711574
LogicalPlanBuilder::from(plan)
1572-
.window(vec![
1573-
Expr::WindowFunction(Box::new(window)).alias(&row_number_name),
1574-
])?
1575+
.window(vec![Expr::WindowFunction(Box::new(window)).alias(name)])?
15751576
.build()
15761577
};
1577-
let left_plan = with_row_number(left_plan, &left_columns)?;
1578-
let right_plan = with_row_number(right_plan, &right_columns)?;
1578+
let left_plan = with_row_number(left_plan, &left_columns, &row_number_name)?;
1579+
let right_plan =
1580+
with_row_number(right_plan, &right_columns, &right_row_number_name)?;
15791581
let (left_builder, right_builder, requalified) = requalify_sides_if_needed(
15801582
LogicalPlanBuilder::from(left_plan),
15811583
LogicalPlanBuilder::from(right_plan),
15821584
)?;
15831585
let left_plan = left_builder.build()?;
15841586
let right_plan = right_builder.build()?;
1585-
let right_columns = right_plan.schema().columns();
1586-
let join_keys = left_plan
1587-
.schema()
1588-
.fields()
1587+
// Requalifying can also rename columns (`x` becomes `x:1` when two
1588+
// relations on one side both have an `x`), so read the join keys from
1589+
// the schemas the join actually sees.
1590+
let left_join_columns = left_plan.schema().columns();
1591+
let right_join_columns = right_plan.schema().columns();
1592+
let join_keys = left_join_columns
15891593
.iter()
1590-
.zip(right_plan.schema().fields().iter())
1591-
.map(|(left_field, right_field)| {
1592-
(
1593-
Column::from_name(left_field.name()),
1594-
Column::from_name(right_field.name()),
1595-
)
1596-
})
1594+
.cloned()
1595+
.zip(right_join_columns.iter().cloned())
15971596
.collect();
1597+
// The output keeps the left input's qualifiers and names, so a query
1598+
// can keep referring to them after the set operation.
15981599
let projection = left_columns
15991600
.into_iter()
1600-
.zip(right_columns)
1601+
.zip(left_join_columns)
1602+
.zip(right_join_columns)
16011603
.zip(from_right)
16021604
.zip(left_plan.schema().fields())
1603-
.map(|(((column, right_column), from_right), field)| {
1604-
let left_column = if requalified {
1605-
Column::new(Some(TableReference::bare("left")), column.name)
1606-
} else {
1607-
column
1608-
};
1609-
if from_right {
1610-
let target_field = Arc::new(
1611-
Field::new(&left_column.name, field.data_type().clone(), false)
1612-
.with_metadata(field.metadata().clone()),
1613-
);
1614-
Expr::Cast(Cast::new_from_field(
1615-
Box::new(Expr::Column(right_column)),
1616-
target_field,
1617-
))
1618-
.alias_qualified(left_column.relation, &left_column.name)
1619-
} else {
1620-
Expr::Column(left_column)
1621-
}
1622-
})
1605+
.map(
1606+
|((((column, left_column), right_column), from_right), field)| {
1607+
let expr = if from_right {
1608+
let target_field = Arc::new(
1609+
Field::new(&column.name, field.data_type().clone(), false)
1610+
.with_metadata(field.metadata().clone()),
1611+
);
1612+
Expr::Cast(Cast::new_from_field(
1613+
Box::new(Expr::Column(right_column)),
1614+
target_field,
1615+
))
1616+
} else if requalified {
1617+
Expr::Column(left_column)
1618+
} else {
1619+
return Expr::Column(column);
1620+
};
1621+
expr.alias_qualified(column.relation, column.name)
1622+
},
1623+
)
16231624
.collect::<Vec<_>>();
16241625
let joined = LogicalPlanBuilder::from(left_plan).join_detailed(
16251626
right_plan,

‎datafusion/substrait/tests/cases/roundtrip_logical_plan.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1659,7 +1659,9 @@ async fn self_referential_intersect_all() -> Result<()> {
16591659
\n WindowAggr: windowExpr=[[row_number() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]\
16601660
\n Filter: data.a < Int64(5)\
16611661
\n TableScan: data projection=[a], partial_filters=[data.a < Int64(5)]",
1662-
true,
1662+
// The plan keeps the `data` qualifier on its output, but Substrait
1663+
// does not carry qualifiers, so the round trip comes back as `left.a`.
1664+
false,
16631665
)
16641666
.await
16651667
}
@@ -1678,7 +1680,9 @@ async fn self_referential_except_all() -> Result<()> {
16781680
\n WindowAggr: windowExpr=[[row_number() PARTITION BY [data.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]\
16791681
\n Filter: data.a < Int64(5)\
16801682
\n TableScan: data projection=[a], partial_filters=[data.a < Int64(5)]",
1681-
true,
1683+
// The plan keeps the `data` qualifier on its output, but Substrait
1684+
// does not carry qualifiers, so the round trip comes back as `left.a`.
1685+
false,
16821686
)
16831687
.await
16841688
}

0 commit comments

Comments
 (0)