gene-bordegaray commented on code in PR #23416:
URL: https://github.com/apache/datafusion/pull/23416#discussion_r3564120959


##########
datafusion/sqllogictest/test_files/range_partitioning.slt:
##########
@@ -682,5 +682,344 @@ ORDER BY range_key, value;
 35 350
 35 350
 
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
+
+##########
+# TEST 19: Window on Range Partition Column
+# Range([range_key]) colocates equal range_key values, so
+# PARTITION BY range_key is satisfied without a hash repartition.
+##########
+
+query TT
+EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY 
value) FROM range_partitioned;
+----
+physical_plan
+01)ProjectionExec: expr=[range_key@0 as range_key, 
sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER 
BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING 
AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY 
[range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] 
RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY 
[range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] 
RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { 
"sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER 
BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING 
AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING 
AND CURRENT ROW], mode=[Sorted]
+03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], 
preserve_partitioning=[true]
+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
+
+query II
+SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM 
range_partitioned ORDER BY range_key;
+----
+1 10
+5 50
+10 100
+15 150
+20 200
+25 250
+30 300
+35 350
+
+
+##########
+# TEST 20: Unbounded-Frame Window on Range Partition Column
+# The unbounded frame makes DataFusion use WindowAggExec instead of
+# BoundedWindowAggExec, which likewise reuses Range partitioning without a
+# hash repartition.
+##########
+
+query TT
+EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY 
value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM 
range_partitioned;
+----
+physical_plan
+01)ProjectionExec: expr=[range_key@0 as range_key, 
sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER 
BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING 
AND UNBOUNDED FOLLOWING@2 as sum(range_partitioned.value) PARTITION BY 
[range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] 
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]
+02)--WindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY 
[range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] 
ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: 
"sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER 
BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING 
AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: 
WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: 
Following(UInt64(NULL)), is_causal: false }]
+03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], 
preserve_partitioning=[true]
+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
+
+query II
+SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS 
BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned 
ORDER BY range_key;
+----
+1 10
+5 50
+10 100
+15 150
+20 200
+25 250
+30 300
+35 350
+
+
+##########
+# TEST 21: Window on Non-Range Column Rehashes
+# Range([range_key]) does not colocate non_range_key values, so
+# PARTITION BY non_range_key still requires a hash repartition.
+##########
+
+query TT
+EXPLAIN SELECT non_range_key, SUM(value) OVER (PARTITION BY non_range_key 
ORDER BY value) FROM range_partitioned;
+----
+physical_plan
+01)ProjectionExec: expr=[non_range_key@0 as non_range_key, 
sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] 
ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED 
PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY 
[range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS 
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY 
[range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS 
LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { 
"sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] 
ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED 
PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED 
PRECEDING AND CURRENT ROW], mode=[Sorted]

Review Comment:
   what do you think about doing one negative case for `WindowAggExec`



##########
datafusion/sqllogictest/test_files/range_partitioning.slt:
##########
@@ -682,5 +682,344 @@ ORDER BY range_key, value;
 35 350
 35 350
 
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
+
+##########

Review Comment:
   I think the coverage is great here, I do want to avoid this file from 
getting too noisy with checking file preserving and subset satisfaction for 
each operator though. What if we did something like test:
     1. Bounded window accepts exact Range(a) (test 19)
     2. General window accepts exact Range(a) (test 20)
     4. Range(a) does not satisfy PARTITION BY b, and a hash is added for 
bounded and unbounded (test 21 and new test)
     5. Safe range subset, Range(a) satisfies PARTITION BY (a, b) when subset 
is enabled (test 22)
     6. Subset disabled, the same query gets a hash exchange (test 24)
     7. No PARTITION BY,  gets one logical window (test 28)
   Then we should add a matrix test for these configurations as a unit / 
integration test in `enforce_distribution.rs` to ensure we are inserting 
repartitions as expeced according to these. I think this is fitting as those 
configs are tightly coupled with that file. I know @2010YOUY01 had some good 
guidance on test, would love your two cents here too 👍 



##########
datafusion/sqllogictest/test_files/range_partitioning.slt:
##########
@@ -682,5 +682,344 @@ ORDER BY range_key, value;
 35 350
 35 350
 
+statement ok
+set datafusion.execution.target_partitions = 4;
+
+statement ok
+set datafusion.optimizer.subset_repartition_threshold = 4;
+
+statement ok
+set datafusion.optimizer.preserve_file_partitions = 0;
+
+
+##########

Review Comment:
   this is a non-blocker for this PR though, something I am taking note of



-- 
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]

Reply via email to