diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 40dbddfbdf9f..1e9dd14a9e21 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -46,6 +46,7 @@ use datafusion_physical_expr::window::StandardWindowExpr; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; @@ -119,17 +120,19 @@ impl WindowTopN { // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.) let (col_idx, limit_n) = extract_window_limit(filter.predicate())?; - // Step 3: Walk through optional ProjectionExec to find BoundedWindowAggExec + // Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec let child = filter.input(); - let (window_exec, proj_between) = find_window_below(child)?; + let (window_exec, intermediates) = find_window_below(child)?; // Step 4: Verify col_idx references a ROW_NUMBER window output column - let input_field_count = window_exec.input().schema().fields().len(); + let window_exec_typed = window_exec.downcast_ref::()?; + let sort_exec = window_exec_typed.input().downcast_ref::()?; + let input_field_count = window_exec_typed.input().schema().fields().len(); if col_idx < input_field_count { return None; // Filter is on an input column, not a window column } let window_expr_idx = col_idx - input_field_count; - let window_exprs = window_exec.window_expr(); + let window_exprs = window_exec_typed.window_expr(); if window_expr_idx >= window_exprs.len() { return None; } @@ -137,8 +140,7 @@ impl WindowTopN { return None; } - // Step 5: Verify child of window is SortExec - let sort_exec = window_exec.input().downcast_ref::()?; + // Step 5: child of window is SortExec (verified above) let sort_child = sort_exec.input(); // Step 6: Determine partition_prefix_len from the window expression @@ -161,28 +163,19 @@ impl WindowTopN { .ok()?; // Step 8: Rebuild window with new child - let new_window = Arc::clone(&child_as_arc(window_exec)) + let mut result = window_exec .with_new_children(vec![Arc::new(partitioned_topk)]) .ok()?; - // Step 9: If ProjectionExec was between Filter and Window, rebuild it - let result = match proj_between { - Some(proj) => Arc::clone(&child_as_arc(proj)) - .with_new_children(vec![new_window]) - .ok()?, - None => new_window, - }; + // Step 9: Rebuild intermediate nodes (ProjectionExec/RepartitionExec) + for node in intermediates.into_iter().rev() { + result = node.with_new_children(vec![result]).ok()?; + } Some(result) } } -/// Helper to get an `Arc` from a reference. -/// We need this because `with_new_children` takes `Arc`. -fn child_as_arc(plan: &T) -> Arc { - Arc::new(plan.clone()) -} - impl PhysicalOptimizerRule for WindowTopN { fn optimize( &self, @@ -303,29 +296,29 @@ fn is_row_number(expr: &Arc) - udf.fun().name() == "row_number" } +type PlanAndIntermediates = (Arc, Vec>); + /// Walk below a plan node looking for a [`BoundedWindowAggExec`]. /// -/// Handles two cases: -/// - Direct child: `FilterExec → BoundedWindowAggExec` -/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec` +/// Handles sequences of `ProjectionExec` and `RepartitionExec`. /// -/// Returns the window exec and an optional `ProjectionExec` in between, -/// or `None` if no `BoundedWindowAggExec` is found within one or two levels. -fn find_window_below( - plan: &Arc, -) -> Option<(&BoundedWindowAggExec, Option<&ProjectionExec>)> { - // Direct child is BoundedWindowAggExec - if let Some(window) = plan.downcast_ref::() { - return Some((window, None)); - } +/// Returns the window exec and a list of intermediate nodes to rebuild, +/// or `None` if no `BoundedWindowAggExec` is found. +fn find_window_below(plan: &Arc) -> Option { + let mut current = Arc::clone(plan); + let mut intermediates = Vec::new(); - // Child is ProjectionExec with BoundedWindowAggExec below - if let Some(proj) = plan.downcast_ref::() { - let proj_child = proj.input(); - if let Some(window) = proj_child.downcast_ref::() { - return Some((window, Some(proj))); + loop { + if current.downcast_ref::().is_some() { + return Some((current, intermediates)); + } else if current.downcast_ref::().is_some() + || current.downcast_ref::().is_some() + { + let next = Arc::clone(current.children().first()?); + intermediates.push(current); + current = next; + } else { + return None; } } - - None } diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index aee9e52568b0..b08cc7957f30 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -311,6 +311,7 @@ impl ExecutionPlan for PartitionedTopKExec { crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( partition_exprs, )]) + .allow_range_satisfaction_for_key_partitioning() } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index cc3d70a1aea2..9b13384e38eb 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -331,6 +331,7 @@ impl ExecutionPlan for BoundedWindowAggExec { } else { vec![Distribution::KeyPartitioned(self.partition_keys().clone())] }) + .allow_range_satisfaction_for_key_partitioning() } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index f1b78ef5c1a7..7611ff52e93e 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -239,6 +239,7 @@ impl ExecutionPlan for WindowAggExec { } else { vec![Distribution::KeyPartitioned(self.partition_keys())] }) + .allow_range_satisfaction_for_key_partitioning() } fn with_new_children( diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt index 5d004ca7fdfa..4ce7a3ebe5eb 100644 --- a/datafusion/sqllogictest/test_files/range_partitioning.slt +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -684,3 +684,333 @@ ORDER BY range_key, value; statement ok reset datafusion.explain.physical_plan_only; + +########## +# TEST 16: PartitionedTopK on Range Partition Column +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies the TopK partition key and avoids repartitioning. +########## + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.range_key, range_partitioned.value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: range_partitioned projection=[range_key, value] +physical_plan +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] +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] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0], order=[value@1 DESC] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query III +SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key; +---- +1 10 1 +5 50 1 +10 100 1 +15 150 1 +20 200 1 +25 250 1 +30 300 1 +35 350 1 + + +########## +# TEST 17: PartitionedTopK on Non-Range Column +# With subset threshold met and preserve-file disabled, partitioning on a non-range +# key cannot reuse Range([range_key]) and requires hash repartitioning. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.non_range_key, range_partitioned.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 AS rn +02)--Filter: 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(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: range_partitioned projection=[non_range_key, value] +physical_plan +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] +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] +03)----PartitionedTopKExec: fetch=1, partition=[non_range_key@0], order=[value@1 DESC] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=csv, has_header=false + +query III +SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY non_range_key; +---- +1 300 1 +2 350 1 + + +########## +# TEST 18: PartitionedTopK Reuses Range Subset Partitioning +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies partitioning by (range_key, non_range_key). +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.range_key, range_partitioned.non_range_key, range_partitioned.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 AS rn +02)--Filter: 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(1) +03)----WindowAggr: windowExpr=[[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]] +04)------TableScan: range_partitioned projection=[range_key, non_range_key, value] +physical_plan +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] +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] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +query IIII +SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key, non_range_key; +---- +1 1 10 1 +5 2 50 1 +10 1 100 1 +15 2 150 1 +20 1 200 1 +25 2 250 1 +30 1 300 1 +35 2 350 1 + + +########## +# TEST 19: Exact Range PartitionedTopK Below Subset Threshold +# Even when subset satisfaction is disabled, exact Range([range_key]) +# satisfies PARTITION BY range_key when repartitioning would not increase +# partition count. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.range_key, range_partitioned.value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: range_partitioned projection=[range_key, value] +physical_plan +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] +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] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0], order=[value@1 DESC] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + + +########## +# TEST 20: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), +# so it should not satisfy the TopK partition key when subset satisfaction is +# disabled. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.range_key, range_partitioned.non_range_key, range_partitioned.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 AS rn +02)--Filter: 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(1) +03)----WindowAggr: windowExpr=[[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]] +04)------TableScan: range_partitioned projection=[range_key, non_range_key, value] +physical_plan +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] +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] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + + +########## +# TEST 21: PartitionedTopK Rehashes Below Subset Threshold +# With subset threshold 5 and only 4 input partitions, planning repartitions +# to increase parallelism instead of reusing Range partitioning. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.range_key, range_partitioned.value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: range_partitioned projection=[range_key, value] +physical_plan +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] +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] +03)----PartitionedTopKExec: fetch=1, partition=[range_key@0], order=[value@1 DESC] +04)------RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + + +########## +# TEST 22: PartitionedTopK Preserves Range When Preserve File Threshold Met +# With preserve-file threshold 1 and 4 input partitions, Range is preserved +# even though target_partitions is 5. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.range_key, range_partitioned.value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: range_partitioned projection=[range_key, value] +physical_plan +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] +02)--RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 +03)----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] +04)------PartitionedTopKExec: fetch=1, partition=[range_key@0], order=[value@1 DESC] +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 23: PartitionedTopK Rehashes When Preserve File Threshold Not Met +# With preserve-file threshold 5 and only 4 input partitions, planning can +# repartition to increase parallelism. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 5; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +logical_plan +01)Projection: range_partitioned.range_key, range_partitioned.value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: range_partitioned projection=[range_key, value] +physical_plan +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] +02)--RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 +03)----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] +04)------PartitionedTopKExec: fetch=1, partition=[range_key@0], order=[value@1 DESC] +05)--------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch_range_partitioning/range_partitioned/part-3.csv]]}, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=csv, has_header=false + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +statement ok +reset datafusion.explain.physical_plan_only; + +statement ok +reset datafusion.optimizer.enable_window_topn; diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index bf9ce26b3553..2f885991e349 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -624,3 +624,48 @@ DROP TABLE window_topn_nulls; # Reset config to default (false) statement ok SET datafusion.optimizer.enable_window_topn = false; + +statement ok +create table t(c1 int, c2 int) as values (1, 2), (3, 4); + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.repartition_windows = false; + +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +query TT +EXPLAIN SELECT * FROM ( + SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2 DESC) as rn + FROM t +) WHERE rn <= 1; +---- +logical_plan +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 +02)--Filter: row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: t projection=[c1, c2] +physical_plan +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] +02)--RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=1, maintains_sort_order=true +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] +04)------PartitionedTopKExec: fetch=1, partition=[c1@0], order=[c2@1 DESC] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_windows = true; + +statement ok +set datafusion.execution.batch_size = 8192; + +statement ok +set datafusion.optimizer.enable_window_topn = false;