Skip to content

Commit ba695d4

Browse files
adriangbclaude
andcommitted
test: cover renamed and swapped columns in leaf projection pushdown
Add the statements from the report as sqllogictest cases: a rename beside a same-name alias under a filter, a limit, an order by and a group by, a swap of two column names, and a swap where the struct field name is also a column name. Add two optimizer unit tests: one for the ambiguous schema, one for the recovery projection that must keep a rename alive. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 parent c67e79a commit ba695d4

2 files changed

Lines changed: 153 additions & 0 deletions

File tree

‎datafusion/optimizer/src/extract_leaf_expressions.rs‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2880,6 +2880,84 @@ mod tests {
28802880
"#)
28812881
}
28822882

2883+
/// A projection that swaps two column names, below a filter and a struct
2884+
/// field read. The merge must not resolve the parent's column references
2885+
/// through the rename map: the input column `test.user` beside the output
2886+
/// field `user` of the other rename gives an ambiguous schema (#25446).
2887+
#[test]
2888+
fn test_extract_above_projection_that_swaps_column_names() -> Result<()> {
2889+
let table_scan = test_table_scan_with_struct()?;
2890+
let plan = LogicalPlanBuilder::from(table_scan)
2891+
.project(vec![col("user").alias("id"), col("id").alias("user")])?
2892+
.filter(col("user").gt(lit(0u32)))?
2893+
.project(vec![col("id"), col("user"), leaf_udf(col("id"), "name")])?
2894+
.build()?;
2895+
2896+
assert_stages!(plan, @r#"
2897+
## Original Plan
2898+
Projection: id, user, leaf_udf(id, Utf8("name"))
2899+
Filter: user > UInt32(0)
2900+
Projection: test.user AS id, test.id AS user
2901+
TableScan: test projection=[id, user]
2902+
2903+
## After Extraction
2904+
(same as original)
2905+
2906+
## After Pushdown
2907+
Projection: id, user, __datafusion_extracted_1 AS leaf_udf(id,Utf8("name"))
2908+
Filter: user > UInt32(0)
2909+
Projection: test.user AS id, test.id AS user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1
2910+
TableScan: test projection=[id, user]
2911+
2912+
## Optimized
2913+
(same as after pushdown)
2914+
"#)
2915+
}
2916+
2917+
/// A projection that gives a computed column the name of its own input
2918+
/// column. The extraction projection goes below the filter, which puts a
2919+
/// column named `id` under the rename, so the two plans have the same set
2920+
/// of field names. The recovery projection must stay, or the rename is lost
2921+
/// and the query gives wrong results.
2922+
#[test]
2923+
fn test_extract_above_projection_that_redefines_column_name() -> Result<()> {
2924+
let table_scan = test_table_scan_with_struct()?;
2925+
let plan = LogicalPlanBuilder::from(table_scan)
2926+
.filter(col("id").gt(lit(0u32)))?
2927+
.project(vec![(col("id") * lit(10u32)).alias("id"), col("user")])?
2928+
.alias("sub")?
2929+
.project(vec![col("sub.id"), leaf_udf(col("sub.user"), "name")])?
2930+
.build()?;
2931+
2932+
assert_stages!(plan, @r#"
2933+
## Original Plan
2934+
Projection: sub.id, leaf_udf(sub.user, Utf8("name"))
2935+
SubqueryAlias: sub
2936+
Projection: test.id * UInt32(10) AS id, test.user
2937+
Filter: test.id > UInt32(0)
2938+
TableScan: test projection=[id, user]
2939+
2940+
## After Extraction
2941+
(same as original)
2942+
2943+
## After Pushdown
2944+
Projection: sub.id, __datafusion_extracted_1 AS leaf_udf(sub.user,Utf8("name"))
2945+
SubqueryAlias: sub
2946+
Projection: test.id * UInt32(10) AS id, test.user, __datafusion_extracted_1
2947+
Filter: test.id > UInt32(0)
2948+
Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.id, test.user
2949+
TableScan: test projection=[id, user]
2950+
2951+
## Optimized
2952+
Projection: sub.id, __datafusion_extracted_1 AS leaf_udf(sub.user,Utf8("name"))
2953+
SubqueryAlias: sub
2954+
Projection: test.id * UInt32(10) AS id, __datafusion_extracted_1
2955+
Filter: test.id > UInt32(0)
2956+
Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1, test.id
2957+
TableScan: test projection=[id, user]
2958+
"#)
2959+
}
2960+
28832961
// =========================================================================
28842962
// SubqueryAlias extraction tests
28852963
// =========================================================================

‎datafusion/sqllogictest/test_files/struct.slt‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1803,3 +1803,78 @@ drop view struct_ctor_view;
18031803

18041804
statement ok
18051805
drop table struct_ctor_null;
1806+
1807+
# A sub-query projection that renames a column to the name of a different
1808+
# column of the same input, with a struct field read above it. Leaf projection
1809+
# pushdown resolved the parent's column references through the rename map and
1810+
# added the renamed input column a second time, which made the output schema
1811+
# ambiguous (https://github.com/apache/datafusion/issues/25446).
1812+
statement ok
1813+
create table rename_swap_struct(a int, b int, s struct<x varchar>) as values (1, 10, {x: 'p'}), (2, 20, {x: 'q'});
1814+
1815+
# a rename beside a same-name alias, under a filter
1816+
query IIT
1817+
select b, a, s['x'] from (select rename_swap_struct.a as b, rename_swap_struct.b as a, s from rename_swap_struct) where a > 0 order by b;
1818+
----
1819+
1 10 p
1820+
2 20 q
1821+
1822+
# the same shape under a limit
1823+
query IIT rowsort
1824+
select b, a, s['x'] from (select rename_swap_struct.a as b, rename_swap_struct.b as a, s from rename_swap_struct) limit 10;
1825+
----
1826+
1 10 p
1827+
2 20 q
1828+
1829+
# a swap of two column names
1830+
query IIT
1831+
select a, c, s['x'] from (select rename_swap_struct.b as a, rename_swap_struct.a as c, s from rename_swap_struct) where a > 0 order by a;
1832+
----
1833+
10 1 p
1834+
20 2 q
1835+
1836+
# the same shape under an order by
1837+
query IIT
1838+
select b, a, s['x'] from (select rename_swap_struct.a as b, rename_swap_struct.b as a, s from rename_swap_struct) order by b;
1839+
----
1840+
1 10 p
1841+
2 20 q
1842+
1843+
# the same shape under a group by
1844+
query III
1845+
select b, a, count(s['x']) from (select rename_swap_struct.a as b, rename_swap_struct.b as a, s from rename_swap_struct) group by b, a order by b;
1846+
----
1847+
1 10 1
1848+
2 20 1
1849+
1850+
# the rename stays above the extraction projection, and the struct field is
1851+
# still read at the scan
1852+
query TT
1853+
explain select b, a, s['x'] from (select rename_swap_struct.a as b, rename_swap_struct.b as a, s from rename_swap_struct) where a > 0;
1854+
----
1855+
logical_plan
1856+
01)Projection: rename_swap_struct.a AS b, rename_swap_struct.b AS a, __datafusion_extracted_1 AS rename_swap_struct.s[x]
1857+
02)--Filter: rename_swap_struct.b > Int32(0)
1858+
03)----Projection: get_field(rename_swap_struct.s, Utf8("x")) AS __datafusion_extracted_1, rename_swap_struct.a, rename_swap_struct.b
1859+
04)------TableScan: rename_swap_struct projection=[a, b, s]
1860+
physical_plan
1861+
01)ProjectionExec: expr=[a@1 as b, b@2 as a, __datafusion_extracted_1@0 as rename_swap_struct.s[x]]
1862+
02)--FilterExec: b@2 > 0
1863+
03)----ProjectionExec: expr=[get_field(s@2, x) as __datafusion_extracted_1, a@0 as a, b@1 as b]
1864+
04)------DataSourceExec: partitions=1, partition_sizes=[1]
1865+
1866+
statement ok
1867+
drop table rename_swap_struct;
1868+
1869+
# The same swap where the struct field name is also a column name of the table.
1870+
statement ok
1871+
create table rename_swap_struct_field(a int, b int, s struct<b varchar>) as values (1, 10, {b: 'p'}), (2, 20, {b: 'q'});
1872+
1873+
query IIT
1874+
select a, b, s['b'] from (select b as a, a as b, s from rename_swap_struct_field limit 100) order by a;
1875+
----
1876+
10 1 p
1877+
20 2 q
1878+
1879+
statement ok
1880+
drop table rename_swap_struct_field;

0 commit comments

Comments
 (0)