[FEAT] Pre Shuffle Merge Strategy

[colin-ho]

Enables an experimental pre-shuffle merge strategy for shuffles.

Our traditional shuffle strategy is a fully materializing map-reduce. Each of the M map partition tasks ends in a fanout, producing N output partitions, and then N reduce partition tasks are ran to merge M outputs each.

The pre-shuffle merge strategy works by merging P map partition tasks, before doing a fanout. The N reduce partition tasks will now only have to merge M/P outputs. The benefit here is that we can reduce the total number of intermediate objects by a factor of P, i.e. from M * N to M / P * N.

I tested this on a 1000 x 1000 shuffle of 100kb partitions. with P = 8

• NaiveMapReduce: 8:44s to complete only 50% of reduces. After which I killed the job because I got tired of waiting.
• PreShuffleMerge: 5:50s to finish the whole workload.

There are 2 big TODOs here:

1. How do we figure out P. We can make this dynamic by setting some memory size threshold, or statically configure it based on number of map tasks and some assumptions on node memory capacity. Even if it's just 2 though, thats 50% reduction in intermediate objects, which can probably provide some decent improvements.
2. Locality. Ideally we want to colocate the maps and the merges.

[codspeed-hq]

CodSpeed Performance Report

Merging #3191 will not alter performance

_{Comparing colin/preshufflemerge (8a494a1) with main (0d669ca)}

Summary

✅ 17 untouched benchmarks

[codecov]

Codecov Report

Attention: Patch coverage is 37.77778% with 168 lines in your changes missing coverage. Please review.

Project coverage is 78.96%. Comparing base (701a011) to head (8a494a1).
Report is 28 commits behind head on main.

Files with missing lines	Patch %	Lines
daft/io/_generated.py	0.00%	34 Missing ⚠️
src/daft-plan/src/physical_ops/shuffle_exchange.rs	38.88%	33 Missing ⚠️
src/daft-scheduler/src/scheduler.rs	0.00%	33 Missing ⚠️
daft/execution/execution_step.py	22.85%	27 Missing ⚠️
daft/execution/physical_plan.py	45.83%	13 Missing ⚠️
src/daft-scan/src/python.rs	52.94%	8 Missing ⚠️
...sical_optimization/rules/reorder_partition_keys.rs	45.45%	6 Missing ⚠️
src/daft-scan/src/lib.rs	25.00%	6 Missing ⚠️
src/common/daft-config/src/python.rs	42.85%	4 Missing ⚠️
src/daft-micropartition/src/python.rs	78.94%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3191      +/-   ##
==========================================
+ Coverage   78.82%   78.96%   +0.14%     
==========================================
  Files         621      638      +17     
  Lines       74922    78196    +3274     
==========================================
+ Hits        59058    61751    +2693     
- Misses      15864    16445     +581     

Files with missing lines	Coverage Δ
daft/context.py	79.75% <ø> (-0.25%)	⬇️
daft/io/_lance.py	92.85% <ø> (ø)
src/common/daft-config/src/lib.rs	87.71% <100.00%> (+0.21%)	⬆️
...rc/physical_optimization/rules/drop_repartition.rs	100.00% <100.00%> (ø)
src/daft-plan/src/physical_planner/translate.rs	95.66% <100.00%> (-0.19%)	⬇️
src/common/daft-config/src/python.rs	76.59% <42.85%> (-0.76%)	⬇️
src/daft-micropartition/src/python.rs	78.02% <78.94%> (-0.21%)	⬇️
...sical_optimization/rules/reorder_partition_keys.rs	77.36% <45.45%> (-2.08%)	⬇️
src/daft-scan/src/lib.rs	63.53% <25.00%> (+0.39%)	⬆️
src/daft-scan/src/python.rs	74.53% <52.94%> (-1.20%)	⬇️
... and 5 more

... and 123 files with indirect coverage changes