feat: Add sort-based shuffle implementation

[andygrove]

Summary

This PR implements a sort-based shuffle mechanism for Ballista, similar to Spark's approach. Instead of creating one file per output partition (N×M files for N input partitions and M output partitions), each input partition writes to a single consolidated file with an index file mapping partition IDs to batch ranges.

Key Features

• Consolidated output files: Each input partition produces one data file + one index file (2×N files instead of N×M)
• Arrow IPC File format: Uses FileWriter/FileReader for efficient random access to specific record batches
• Memory management: Configurable buffer sizes with spill-to-disk support to prevent OOM
• Backward compatible: Sort-based shuffle is opt-in via configuration; default remains hash-based shuffle

Architecture

ballista/core/src/execution_plans/sort_shuffle/
├── mod.rs          # Module exports
├── config.rs       # SortShuffleConfig with buffer/memory settings
├── buffer.rs       # PartitionBuffer for in-memory buffering
├── spill.rs        # SpillManager for disk spilling
├── index.rs        # ShuffleIndex for batch range tracking
├── writer.rs       # SortShuffleWriterExec execution plan
└── reader.rs       # Reading logic with FileReader random access

Configuration Options

• ballista.shuffle.sort_based.enabled - Enable sort-based shuffle (default: false)
• ballista.shuffle.sort_based.buffer_size - Per-partition buffer size in bytes (default: 1MB)
• ballista.shuffle.sort_based.memory_limit - Total memory limit for buffers (default: 256MB)

Performance

The sort-based shuffle can be more efficient when:

• There are many output partitions (reduces file count from N×M to 2×N)
• Network overhead of many small files is significant
• Sequential I/O patterns are preferred over random I/O

Benchmark (cargo run --release --bin shuffle_bench):

• Hash shuffle: Creates M files per input partition
• Sort shuffle: Creates 1 data + 1 index file per input partition
• FileReader random access: Directly seeks to needed batches without scanning

Testing

• Unit tests for each component (buffer, spill, index, reader, writer)
• End-to-end integration tests (ballista/client/tests/sort_shuffle.rs)
• Comparison tests verifying sort shuffle produces same results as hash shuffle

Changes

Core

• Added sort_shuffle module with all components
• Added ShuffleWriter trait for polymorphic shuffle writer handling
• Updated shuffle_reader.rs to detect and read sort-based shuffle format

Executor

• Updated DefaultExecutionEngine to handle both ShuffleWriterExec and SortShuffleWriterExec
• Added ShuffleWriterVariant enum for type-safe dispatch

Scheduler/Planner

• Updated DistributedPlanner to create SortShuffleWriterExec when enabled
• Updated ExecutionStageBuilder to use ShuffleWriter trait

Benchmarks

• Added shuffle_bench binary for comparing hash vs sort shuffle performance

Related Issues

This PR relates to several open issues in apache/datafusion-ballista:

• [EPIC] Shuffle file execs improvement apache/datafusion-ballista#1319 - [EPIC] Shuffle file execs improvement
• Investigate use of comet shuffle apache/datafusion-ballista#1287 - Investigate use of comet shuffle
• Implement memory management in shuffle writer apache/datafusion-ballista#320 - Implement memory management in shuffle writer
• [Ballista] Streaming style push-based shuffle and All-at-once stage scheduling in Ballista apache/datafusion-ballista#1151 - Streaming style push-based shuffle and All-at-once stage scheduling

---

🤖 Generated with Claude Code