Skip to content

Commit e98fb60

Browse files
stuhoodGene Bordegaray
stuhood
authored and
Gene Bordegaray
committed
Allow Range partitioned inputs to PartitionedTopK (apache#23355)
- Closes apache#23290. In two commits: * Improve the robustness of `WindowTopN` to intermediate `RepartitionExec` nodes, to ensure that `PartitionedTopK` can still be used when `target_partitions` mismatches physical partitions. * Use `required_input_distributions` to declare required partitioning for `PartitionedTopKExec`, `BoundedWindowAggExec`, and `WindowAggExec`. Yes, new `slt` tests are added for both the `WindowTopN` robustness fix, and for the use of `required_input_distributions` in `PartitionedTopK`. No: only plan changes.
1 parent 1248881 commit e98fb60

4 files changed

Lines changed: 238 additions & 39 deletions

File tree

‎datafusion/physical-optimizer/src/window_topn.rs‎

Lines changed: 35 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ use datafusion_physical_expr::window::StandardWindowExpr;
4646
use datafusion_physical_plan::ExecutionPlan;
4747
use datafusion_physical_plan::filter::FilterExec;
4848
use datafusion_physical_plan::projection::ProjectionExec;
49+
use datafusion_physical_plan::repartition::RepartitionExec;
4950
use datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec;
5051
use datafusion_physical_plan::sorts::sort::SortExec;
5152
use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
@@ -119,26 +120,27 @@ impl WindowTopN {
119120
// Step 2: Extract limit from predicate (rn <= K, rn < K, etc.)
120121
let (col_idx, limit_n) = extract_window_limit(filter.predicate())?;
121122

122-
// Step 3: Walk through optional ProjectionExec to find BoundedWindowAggExec
123+
// Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec
123124
let child = filter.input();
124-
let (window_exec, proj_between) = find_window_below(child)?;
125+
let (window_exec, intermediates) = find_window_below(child)?;
125126

126127
// Step 4: Verify col_idx references a ROW_NUMBER window output column
127-
let input_field_count = window_exec.input().schema().fields().len();
128+
let window_exec_typed = window_exec.downcast_ref::<BoundedWindowAggExec>()?;
129+
let sort_exec = window_exec_typed.input().downcast_ref::<SortExec>()?;
130+
let input_field_count = window_exec_typed.input().schema().fields().len();
128131
if col_idx < input_field_count {
129132
return None; // Filter is on an input column, not a window column
130133
}
131134
let window_expr_idx = col_idx - input_field_count;
132-
let window_exprs = window_exec.window_expr();
135+
let window_exprs = window_exec_typed.window_expr();
133136
if window_expr_idx >= window_exprs.len() {
134137
return None;
135138
}
136139
if !is_row_number(&window_exprs[window_expr_idx]) {
137140
return None;
138141
}
139142

140-
// Step 5: Verify child of window is SortExec
141-
let sort_exec = window_exec.input().downcast_ref::<SortExec>()?;
143+
// Step 5: child of window is SortExec (verified above)
142144
let sort_child = sort_exec.input();
143145

144146
// Step 6: Determine partition_prefix_len from the window expression
@@ -161,28 +163,19 @@ impl WindowTopN {
161163
.ok()?;
162164

163165
// Step 8: Rebuild window with new child
164-
let new_window = Arc::clone(&child_as_arc(window_exec))
166+
let mut result = window_exec
165167
.with_new_children(vec![Arc::new(partitioned_topk)])
166168
.ok()?;
167169

168-
// Step 9: If ProjectionExec was between Filter and Window, rebuild it
169-
let result = match proj_between {
170-
Some(proj) => Arc::clone(&child_as_arc(proj))
171-
.with_new_children(vec![new_window])
172-
.ok()?,
173-
None => new_window,
174-
};
170+
// Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec)
171+
for node in intermediates.into_iter().rev() {
172+
result = node.with_new_children(vec![result]).ok()?;
173+
}
175174

176175
Some(result)
177176
}
178177
}
179178

180-
/// Helper to get an `Arc<dyn ExecutionPlan>` from a reference.
181-
/// We need this because `with_new_children` takes `Arc<Self>`.
182-
fn child_as_arc<T: ExecutionPlan + Clone + 'static>(plan: &T) -> Arc<dyn ExecutionPlan> {
183-
Arc::new(plan.clone())
184-
}
185-
186179
impl PhysicalOptimizerRule for WindowTopN {
187180
fn optimize(
188181
&self,
@@ -303,29 +296,32 @@ fn is_row_number(expr: &Arc<dyn datafusion_physical_expr::window::WindowExpr>) -
303296
udf.fun().name() == "row_number"
304297
}
305298

299+
type PlanAndIntermediates = (Arc<dyn ExecutionPlan>, Vec<Arc<dyn ExecutionPlan>>);
300+
306301
/// Walk below a plan node looking for a [`BoundedWindowAggExec`].
307302
///
308-
/// Handles two cases:
309-
/// - Direct child: `FilterExec → BoundedWindowAggExec`
310-
/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec`
303+
/// Handles sequences of `ProjectionExec` and `RepartitionExec`.
304+
/// This is safe because `PartitionedTopKExec` can be pushed below them:
305+
/// projections only provide aliases, and pushing the limit below repartitions
306+
/// is safe because the limit is computed per-partition.
311307
///
312-
/// Returns the window exec and an optional `ProjectionExec` in between,
313-
/// or `None` if no `BoundedWindowAggExec` is found within one or two levels.
314-
fn find_window_below(
315-
plan: &Arc<dyn ExecutionPlan>,
316-
) -> Option<(&BoundedWindowAggExec, Option<&ProjectionExec>)> {
317-
// Direct child is BoundedWindowAggExec
318-
if let Some(window) = plan.downcast_ref::<BoundedWindowAggExec>() {
319-
return Some((window, None));
320-
}
308+
/// Returns the window exec and a list of intermediate nodes to rebuild,
309+
/// or `None` if no `BoundedWindowAggExec` is found.
310+
fn find_window_below(plan: &Arc<dyn ExecutionPlan>) -> Option<PlanAndIntermediates> {
311+
let mut current = Arc::clone(plan);
312+
let mut intermediates = Vec::new();
321313

322-
// Child is ProjectionExec with BoundedWindowAggExec below
323-
if let Some(proj) = plan.downcast_ref::<ProjectionExec>() {
324-
let proj_child = proj.input();
325-
if let Some(window) = proj_child.downcast_ref::<BoundedWindowAggExec>() {
326-
return Some((window, Some(proj)));
314+
loop {
315+
if current.downcast_ref::<BoundedWindowAggExec>().is_some() {
316+
return Some((current, intermediates));
317+
} else if current.downcast_ref::<ProjectionExec>().is_some()
318+
|| current.downcast_ref::<RepartitionExec>().is_some()
319+
{
320+
let next = Arc::clone(current.children().first()?);
321+
intermediates.push(current);
322+
current = next;
323+
} else {
324+
return None;
327325
}
328326
}
329-
330-
None
331327
}

‎datafusion/physical-plan/src/sorts/partitioned_topk.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,7 @@ impl ExecutionPlan for PartitionedTopKExec {
314314
crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned(
315315
partition_exprs,
316316
)])
317+
.allow_range_satisfaction_for_key_partitioning()
317318
}
318319

319320
fn maintains_input_order(&self) -> Vec<bool> {

‎datafusion/sqllogictest/test_files/range_partitioning.slt‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -897,5 +897,162 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER
897897
30 1060
898898
35 1410
899899

900+
901+
902+
##########
903+
# TEST 26: PartitionedTopK on Range Partition Column
904+
# Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning.
905+
##########
906+
907+
statement ok
908+
set datafusion.optimizer.enable_window_topn = true;
909+
910+
statement ok
911+
set datafusion.optimizer.subset_repartition_threshold = 4;
912+
913+
statement ok
914+
set datafusion.optimizer.preserve_file_partitions = 0;
915+
916+
query TT
917+
EXPLAIN SELECT * FROM (
918+
SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn
919+
FROM range_partitioned
920+
) WHERE rn <= 1;
921+
----
922+
physical_plan
923+
01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn]
924+
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
925+
03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0], order=[value@1 DESC]
926+
04)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
927+
928+
query III
929+
SELECT * FROM (
930+
SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn
931+
FROM range_partitioned
932+
) WHERE rn <= 1
933+
ORDER BY range_key;
934+
----
935+
1 10 1
936+
5 50 1
937+
10 100 1
938+
15 150 1
939+
20 200 1
940+
25 250 1
941+
30 300 1
942+
35 350 1
943+
944+
945+
##########
946+
# TEST 27: PartitionedTopK on Non-Range Column
947+
# Partitioning on a non-range key cannot reuse Range([range_key]) and
948+
# requires hash repartitioning.
949+
##########
950+
951+
statement ok
952+
set datafusion.optimizer.subset_repartition_threshold = 4;
953+
954+
statement ok
955+
set datafusion.optimizer.preserve_file_partitions = 0;
956+
957+
query TT
958+
EXPLAIN SELECT * FROM (
959+
SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn
960+
FROM range_partitioned
961+
) WHERE rn <= 1;
962+
----
963+
physical_plan
964+
01)ProjectionExec: expr=[non_range_key@0 as non_range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn]
965+
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
966+
03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[non_range_key@0], order=[value@1 DESC]
967+
04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4
968+
05)--------DataSourceExec: file_groups=<slt:ignore>, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false
969+
970+
query III
971+
SELECT * FROM (
972+
SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn
973+
FROM range_partitioned
974+
) WHERE rn <= 1
975+
ORDER BY non_range_key;
976+
----
977+
1 300 1
978+
2 350 1
979+
980+
981+
##########
982+
# TEST 28: PartitionedTopK Reuses Range Subset Partitioning
983+
# With subset threshold met and preserve-file disabled, Range([range_key])
984+
# satisfies partitioning by (range_key, non_range_key).
985+
##########
986+
987+
statement ok
988+
set datafusion.optimizer.subset_repartition_threshold = 4;
989+
990+
statement ok
991+
set datafusion.optimizer.preserve_file_partitions = 0;
992+
993+
query TT
994+
EXPLAIN SELECT * FROM (
995+
SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn
996+
FROM range_partitioned
997+
) WHERE rn <= 1;
998+
----
999+
physical_plan
1000+
01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
1001+
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
1002+
03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC]
1003+
04)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
1004+
1005+
query IIII
1006+
SELECT * FROM (
1007+
SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn
1008+
FROM range_partitioned
1009+
) WHERE rn <= 1
1010+
ORDER BY range_key, non_range_key;
1011+
----
1012+
1 1 10 1
1013+
5 2 50 1
1014+
10 1 100 1
1015+
15 2 150 1
1016+
20 1 200 1
1017+
25 2 250 1
1018+
30 1 300 1
1019+
35 2 350 1
1020+
1021+
1022+
##########
1023+
# TEST 29: Range Subset PartitionedTopK Rehashes Below Subset Threshold
1024+
# Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key),
1025+
# so it should not satisfy the TopK partition key when subset satisfaction is
1026+
# disabled.
1027+
##########
1028+
1029+
statement ok
1030+
set datafusion.execution.target_partitions = 4;
1031+
1032+
statement ok
1033+
set datafusion.optimizer.subset_repartition_threshold = 5;
1034+
1035+
statement ok
1036+
set datafusion.optimizer.preserve_file_partitions = 0;
1037+
1038+
query TT
1039+
EXPLAIN SELECT * FROM (
1040+
SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn
1041+
FROM range_partitioned
1042+
) WHERE rn <= 1;
1043+
----
1044+
physical_plan
1045+
01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
1046+
02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
1047+
03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC]
1048+
04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4
1049+
05)--------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false
1050+
1051+
statement ok
1052+
reset datafusion.optimizer.subset_repartition_threshold;
1053+
9001054
statement ok
9011055
reset datafusion.explain.physical_plan_only;
1056+
1057+
statement ok
1058+
reset datafusion.optimizer.enable_window_topn;

‎datafusion/sqllogictest/test_files/window_topn.slt‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -624,3 +624,48 @@ DROP TABLE window_topn_nulls;
624624
# Reset config to default (false)
625625
statement ok
626626
SET datafusion.optimizer.enable_window_topn = false;
627+
628+
statement ok
629+
create table t(c1 int, c2 int) as values (1, 2), (3, 4);
630+
631+
statement ok
632+
set datafusion.execution.target_partitions = 5;
633+
634+
statement ok
635+
set datafusion.optimizer.repartition_windows = false;
636+
637+
statement ok
638+
set datafusion.execution.batch_size = 1;
639+
640+
statement ok
641+
set datafusion.optimizer.enable_window_topn = true;
642+
643+
query TT
644+
EXPLAIN SELECT * FROM (
645+
SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2 DESC) as rn
646+
FROM t
647+
) WHERE rn <= 1;
648+
----
649+
logical_plan
650+
01)Projection: t.c1, t.c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn
651+
02)--Filter: row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1)
652+
03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]]
653+
04)------TableScan: t projection=[c1, c2]
654+
physical_plan
655+
01)ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn]
656+
02)--RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true
657+
03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
658+
04)------PartitionedTopKExec: fn=row_number, fetch=1, partition=[c1@0], order=[c2@1 DESC]
659+
05)--------DataSourceExec: partitions=1, partition_sizes=[1]
660+
661+
statement ok
662+
set datafusion.execution.target_partitions = 4;
663+
664+
statement ok
665+
set datafusion.optimizer.repartition_windows = true;
666+
667+
statement ok
668+
set datafusion.execution.batch_size = 8192;
669+
670+
statement ok
671+
set datafusion.optimizer.enable_window_topn = false;

0 commit comments

Comments
 (0)