Arrow Flight for Distributed Shuffles

[srilman]

Currently, Daft's distributed runner (Flotilla) utilizes Ray object store for peer-to-peer communication of tabular data (aka MicroPartitions). Although Ray's Object Store is readily available when working on Ray clusters and provides some additional benefits (lineage, ref-counting / garbage collection), we have started to notice some issues. In particular:

• We have seen cases where there can be a lot of overhead / contention when interacting with the object store. A part of it seems to be pickling overhead.
• Lineage tracking requires a large amount of memory per object, which is all stored on the head node. We have seen the head node OOMing when performing 10TB+ shuffles.
• We can't always implement additional optimizations on top, such as async & streaming IO, zero-copy serialization, etc.

Thus, in addition to our Ray-based shuffle service, we are adding a native shuffle service called Flight Shuffle that:

• Based on the Arrow Flight RPC protocol, communicates via the Arrow format
• Peer-to-peer communication
• Designed for local NVMe storage
• Built-in scheduler fault-tolerance & task lineage

We can break this down into 3 sub-tasks.

1) Basic Flight Shuffle

Build and tune Flight shuffle to replace our Ray-based repartition (used for GroupBy and Join operations) such that the performance is within 20% of Ray's version.

Tasks:

• Implement a basic version of Flight shuffle
• Profile against various benchmarks and iteratively tune
  • TPCH
  • MinHash Dedupe
  • TPC-DS?
  • (TBD)

Any potential follow-up tasks or optimizations will be listed below.

2) "Genericize" Shuffle Backends

Rather than hard-coding special support for Ray and Flight shuffles, we should instead track shuffle "backends" as generic trait that supports writing MicroPartitions to and reading from. For more details on the specific design and interface, see the design doc.

Tasks:

• Implement the generic shuffle interface and apply to Flight shuffle
• Apply the generic shuffle interface to Ray shuffle for repartitions
• Extend the interface usage across all shuffling operators
  • Gather
  • Pre-Shuffle Merge
  • Into Partitions
  • Into Batches
  • Sort
  • Limit
• (Optional) Add a basic distributed filesystem-based shuffle

3) Task Lineage and Fault Tolerance

Right now, Daft can only handle worker / task failure for map-only pipelines, where all inputs are hard-coded external items (such as a file path on S3). To support general pipelines, the scheduler would need to determine if intermediate partitions are available, and if not, recursively re-run previous tasks to regenerate them. For more details, see the design doc.

Tasks:

• Track existence of intermediate partitions (and their associated task) in scheduler
• Update partition status with worker deaths and task failures due to missing inputs
• Retry sub-task if input if missing and needed
• Rework pre-shuffle merge with lineage support (TBD)

[chenghuichen]

Picked up a couple of small items here — #6808 and #6809. Happy to adjust if these don't align with your plans.