Skip to content

Commit 0ca8891

Browse files
committed
test: WindowTopN tie-order-independent slt checks, reordered-projection negative test
1 parent f120e98 commit 0ca8891

2 files changed

Lines changed: 46 additions & 31 deletions

File tree

‎datafusion/core/tests/physical_optimizer/window_topn.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -565,6 +565,18 @@ fn repartition_between_passes_the_index_through() -> Result<()> {
565565
Ok(())
566566
}
567567

568+
/// A projection that reorders the window's columns: the filter's `pk@2` maps
569+
/// back to `pk@0`, an input column, so the rule leaves the plan alone. Read as
570+
/// an index into the window's output, it would be `rn`.
571+
#[test]
572+
fn reordered_projection_input_column_no_change() -> Result<()> {
573+
let projection = project(build_rn_rk_window()?, &[2, 3, 0])?; // [rn, rk, pk]
574+
let plan = filter_lteq_3(projection, 2)?;
575+
let before = plan_str(plan.as_ref());
576+
assert_eq!(plan_str(optimize(plan)?.as_ref()), before);
577+
Ok(())
578+
}
579+
568580
/// A column the projection computes has no single window column behind it,
569581
/// so the rule leaves the plan alone. `x` sits at index 2 so that treating it
570582
/// as a pass-through would land on `rn` rather than be rejected as an input

‎datafusion/sqllogictest/test_files/window_topn.slt‎

Lines changed: 34 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1556,19 +1556,21 @@ SET datafusion.optimizer.enable_window_topn = true;
15561556
statement ok
15571557
SET datafusion.explain.physical_plan_only = true;
15581558

1559-
# Filter on rank(): both tied rows per partition are kept.
1560-
query IIII
1561-
SELECT pk, id, rn, rk FROM (
1562-
SELECT pk, v, id,
1563-
row_number() OVER (PARTITION BY pk ORDER BY v) AS rn,
1564-
rank() OVER (PARTITION BY pk ORDER BY v) AS rk
1565-
FROM window_topn_shift
1566-
) WHERE rk <= 1 ORDER BY pk, id;
1567-
----
1568-
1 1 1 1
1569-
1 2 2 1
1570-
2 4 1 1
1571-
2 5 2 1
1559+
# Filter on rank(): both tied rows per partition are kept. `ORDER BY v` leaves
1560+
# the row numbers of tied rows unspecified, so ids and row numbers are checked
1561+
# as sets rather than paired up.
1562+
query I??
1563+
SELECT pk, array_agg(id ORDER BY id), array_agg(rn ORDER BY rn) FROM (
1564+
SELECT pk, id, rn, rk FROM (
1565+
SELECT pk, v, id,
1566+
row_number() OVER (PARTITION BY pk ORDER BY v) AS rn,
1567+
rank() OVER (PARTITION BY pk ORDER BY v) AS rk
1568+
FROM window_topn_shift
1569+
) WHERE rk <= 1
1570+
) GROUP BY pk ORDER BY pk;
1571+
----
1572+
1 [1, 2] [1, 2]
1573+
2 [4, 5] [1, 2]
15721574

15731575
query TT
15741576
EXPLAIN SELECT pk, id, rn, rk FROM (
@@ -1586,17 +1588,18 @@ physical_plan
15861588
05)--------DataSourceExec: partitions=1, partition_sizes=[1]
15871589

15881590
# Filter on row_number(), with the window computing rank() first: no row with
1589-
# rn = 2 comes back.
1590-
query IIII
1591-
SELECT pk, id, rk, rn FROM (
1591+
# rn = 2 comes back. Either tied row may be the one kept; ids 3 and 6 are the
1592+
# only rows with v = 1.
1593+
query IIIB
1594+
SELECT pk, rk, rn, id NOT IN (3, 6) FROM (
15921595
SELECT pk, v, id,
15931596
rank() OVER (PARTITION BY pk ORDER BY v) AS rk,
15941597
row_number() OVER (PARTITION BY pk ORDER BY v) AS rn
15951598
FROM window_topn_shift
1596-
) WHERE rn <= 1 ORDER BY pk, id;
1599+
) WHERE rn <= 1 ORDER BY pk;
15971600
----
1598-
1 1 1 1
1599-
2 4 1 1
1601+
1 1 1 true
1602+
2 1 1 true
16001603

16011604
query TT
16021605
EXPLAIN SELECT pk, id, rk, rn FROM (
@@ -1614,18 +1617,18 @@ physical_plan
16141617
05)--------DataSourceExec: partitions=1, partition_sizes=[1]
16151618

16161619
# Filter on dense_rank(), behind row_number().
1617-
query IIII
1618-
SELECT pk, id, rn, dr FROM (
1619-
SELECT pk, v, id,
1620-
row_number() OVER (PARTITION BY pk ORDER BY v) AS rn,
1621-
dense_rank() OVER (PARTITION BY pk ORDER BY v) AS dr
1622-
FROM window_topn_shift
1623-
) WHERE dr <= 1 ORDER BY pk, id;
1624-
----
1625-
1 1 1 1
1626-
1 2 2 1
1627-
2 4 1 1
1628-
2 5 2 1
1620+
query I??
1621+
SELECT pk, array_agg(id ORDER BY id), array_agg(rn ORDER BY rn) FROM (
1622+
SELECT pk, id, rn, dr FROM (
1623+
SELECT pk, v, id,
1624+
row_number() OVER (PARTITION BY pk ORDER BY v) AS rn,
1625+
dense_rank() OVER (PARTITION BY pk ORDER BY v) AS dr
1626+
FROM window_topn_shift
1627+
) WHERE dr <= 1
1628+
) GROUP BY pk ORDER BY pk;
1629+
----
1630+
1 [1, 2] [1, 2]
1631+
2 [4, 5] [1, 2]
16291632

16301633
query TT
16311634
EXPLAIN SELECT pk, id, rn, dr FROM (

0 commit comments

Comments
 (0)