Arrow Flight Shuffle for Comet

[Shekharrajak]

What is the problem the feature request solves?

Replace disk-based shuffle with Arrow Flight for direct memory-to-memory data exchange between executors, eliminating intermediate disk I/O and leveraging Arrow's native IPC
format for efficient shuffle.

Motivation

Comet already uses Arrow RecordBatches internally, but shuffle still goes through disk:

Current Flow:
Arrow RecordBatch → Arrow IPC → Compress → DISK → Network → DISK → Decompress → Arrow RecordBatch

Proposed Flow:
Arrow RecordBatch → Arrow Flight (gRPC) → Arrow RecordBatch

Describe the potential solution

No response

Additional context

Configuration

  # Enable Arrow Flight shuffle
  spark.shuffle.manager=org.apache.comet.shuffle.CometFlightShuffleManager

  # Flight server configuration
  spark.comet.shuffle.flight.enabled=true
  spark.comet.shuffle.flight.port=50051

  # Memory management
  spark.comet.shuffle.flight.memoryFraction=0.3
  spark.comet.shuffle.flight.spillThreshold=0.8

  # Network configuration
  spark.comet.shuffle.flight.maxMessageSize=67108864  # 64MB
  spark.comet.shuffle.flight.compression=zstd

  # Fault tolerance
  spark.comet.shuffle.flight.retryAttempts=3
  spark.comet.shuffle.flight.retryDelayMs=1000

Related Work

• Ballista: DataFusion's distributed query engine uses Arrow Flight
• Dask: Exploring Arrow Flight for task communication
• Ray: Uses gRPC for object transfer (similar concept)
• Spark 3.2 Push-Based Shuffle: Inspiration for push model

References

• https://arrow.apache.org/docs/format/Flight.html
• https://arrow.apache.org/docs/format/FlightSql.html
• https://github.com/apache/datafusion-comet/tree/main/native/core/src/execution/shuffle
• https://spark.apache.org/docs/latest/api/java/org/apache/spark/shuffle/ShuffleManager.html

[comphead]

Thanks @Shekharrajak it actually sounds attractive, if talking about just replacing transport.
However IMO the flow brings up some challenges:

• less resilient, without disk it would be easier to run out of memory
• backpressure, so if the reducer is busy now, what mapper is supposed to do? buffer in memory? resend? if so we stumble on problem on step1

nevertheless it would be nice having separate shuffle manager for lightweight jobs

[wForget]

Thanks @Shekharrajak it actually sounds attractive, if talking about just replacing transport. However IMO the flow brings up some challenges:

• less resilient, without disk it would be easier to run out of memory
• backpressure, so if the reducer is busy now, what mapper is supposed to do? buffer in memory? resend? if so we stumble on problem on step1

And this may break the Spark DRA feature.

nevertheless it would be nice having separate shuffle manager for lightweight jobs

Introducing a spark remote shuffle service like Apache Celeborn/Uniffle is a good option.