gene-bordegaray commented on code in PR #23355:
URL: https://github.com/apache/datafusion/pull/23355#discussion_r3582159657
##########
datafusion/physical-optimizer/src/window_topn.rs:
##########
@@ -336,29 +329,29 @@ fn supported_window_fn(
}
}
+type PlanAndIntermediates = (Arc<dyn ExecutionPlan>, Vec<Arc<dyn
ExecutionPlan>>);
+
/// Walk below a plan node looking for a [`BoundedWindowAggExec`].
///
-/// Handles two cases:
-/// - Direct child: `FilterExec → BoundedWindowAggExec`
-/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec`
Review Comment:
nit: small comment why projections and repartitions are safe to walk through
##########
datafusion/sqllogictest/test_files/range_partitioning.slt:
##########
@@ -907,5 +907,162 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM
range_partitioned ORDER
30 1060
35 1410
+
+
+##########
+# TEST 30: PartitionedTopK on Range Partition Column
+# Exact 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;
+----
+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: fn=row_number, 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
Review Comment:
can we ignore the file sections 👍
##########
datafusion/sqllogictest/test_files/range_partitioning.slt:
##########
@@ -907,5 +907,162 @@ SELECT range_key, SUM(value) OVER (ORDER BY value) FROM
range_partitioned ORDER
30 1060
35 1410
+
+
+##########
+# TEST 30: PartitionedTopK on Range Partition Column
Review Comment:
I think test numbering is a bit off as well 👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
