refactor(shuffle): Introduce PartitionWriter interface to decouple shuffle partitioning from storage

[wForget]

Which issue does this PR close?

Closes #4780.

Rationale for this change

Comet's shuffle partitioners currently own too much of the local shuffle write
logic directly, including writing shuffle data files, writing index files,
handling spill files, and finalizing partition offsets. This tightly couples the
partitioning logic to the local file-based storage implementation and makes it
hard to introduce alternative shuffle storage backends (e.g. a remote shuffle
writer), since each partitioner would need backend-specific write behavior.

This PR separates shuffle partitioning from shuffle storage so partitioners only
produce partitioned RecordBatch streams, while a storage backend owns how those
batches are persisted and finalized. This creates a clear extension point for
future remote shuffle support without changing the partitioners.

What changes are included in this PR?

• Add a PartitionWriter trait (writers/partition_writer.rs) that abstracts the
  storage backend for shuffle output. Partitioners drive it as: any number of
  write calls to stage batches, one finish_partition per partition in
  ascending id order, then a single finish_all.
• Add LocalPartitionWriter (writers/local/local_partition_writer.rs) that
  implements the existing local file-based behavior behind the new interface:
  • Single-partition mode streams all batches through one long-lived
    BufBatchWriter (preserving cross-batch coalescing).
  • Multi-partition mode stages batches per partition via SpillWriters and
    merges them into the data file one partition at a time.
  • Owns data file writes, index file / partition offset generation, and spill
    handling.
• Move spill.rs under writers/local/ since spilling is a local-storage detail.
• Simplify the three partitioners (SinglePartitionShufflePartitioner,
  MultiPartitionShuffleRepartitioner, EmptySchemaShufflePartitioner) to depend
  only on a PartitionWriter instead of owning file/index/spill logic.
• Construct the LocalPartitionWriter in external_shuffle and inject it into the
  selected partitioner.

This is a refactor only; it preserves the existing local shuffle behavior for
single-partition, multi-partition, and empty-schema shuffles.

How are these changes tested?

Covered by existing tests.

[Copilot]

Pull request overview

This PR refactors Comet’s native shuffle write path to introduce a PartitionWriter abstraction, moving local shuffle file/index writing and spill handling out of the shuffle partitioners. This aligns with issue #4780’s goal of decoupling partitioning logic from the storage backend so alternative backends (like remote shuffle) can be added later.

Changes:

• Introduces a new PartitionWriter trait and a LocalPartitionWriter implementation, plus a renamed SpillWriter for local spill files.
• Updates shuffle partitioners (single_partition, multi_partition, empty_schema) and shuffle_writer wiring to write via the new writer interface.
• Refactors PartitionedBatchIterator to implement Iterator (but currently drops interleave_time metric tracking).

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
native/shuffle/src/writers/rss/mod.rs	Adds placeholder module for future RSS writer implementation.
native/shuffle/src/writers/partition_writer.rs	Adds the new PartitionWriter trait interface.
native/shuffle/src/writers/mod.rs	Reorganizes writer modules and exports to include local and partition_writer.
native/shuffle/src/writers/local/mod.rs	Introduces the local writer module structure.
native/shuffle/src/writers/local/local_partition_writer.rs	Implements LocalPartitionWriter that writes shuffle data + index and manages per-partition spill writers.
native/shuffle/src/writers/local/spill.rs	Refactors old per-partition writer into SpillWriter for spill file handling.
native/shuffle/src/writers/buf_batch_writer.rs	Removes the writer_stream_position helper method.
native/shuffle/src/shuffle_writer.rs	Wires shuffle execution to construct a LocalPartitionWriter and pass it into partitioners; updates related tests.
native/shuffle/src/partitioners/single_partition.rs	Refactors single-partition shuffle to write via PartitionWriter.
native/shuffle/src/partitioners/multi_partition.rs	Refactors multi-partition shuffle and spill to write via PartitionWriter.
native/shuffle/src/partitioners/empty_schema.rs	Refactors empty-schema shuffle to write via PartitionWriter.
native/shuffle/src/partitioners/partitioned_batch_iterator.rs	Makes PartitionedBatchIterator implement Iterator, removing interleave timing from next().
native/shuffle/src/partitioners/mod.rs	Removes re-export of PartitionedBatchIterator.
native/shuffle/src/metrics.rs	Temporarily comments out interleave_time metric field and registration.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.

[andygrove]

Thanks for this, the separation of partitioning from storage reads cleanly and the behavior looks preserved across all three modes.

One design question on the PartitionWriter trait. The generic I: Iterator parameter on write and finish_partition makes the trait non-object-safe, which is why the partitioners all had to become generic over T: PartitionWriter and external_shuffle is pinned to the concrete LocalPartitionWriter. For the remote-backend goal that motivates this refactor, that means adding a second backend would monomorphize each partitioner per backend and external_shuffle would need to branch on backend type while constructing differently-typed partitioners.

Did you consider taking the iterator as &mut dyn Iterator<Item = datafusion::common::Result<RecordBatch>> instead? That would make PartitionWriter object-safe, so the partitioners could hold a Box<dyn PartitionWriter> and external_shuffle could select the backend once and pass it into any partitioner. Curious how you are picturing the remote writer getting wired in from external_shuffle.

Disclosure: I used an LLM to help review this PR and draft this comment.

[wForget]

@andygrove Thanks for your review. Since external_shuffle only needs to use the repartitioner through Box<dyn ShufflePartitioner>, we can pass different PartitionWriter implementations when constructing the ShufflePartitioner and simply return a Result<Box<dyn ShufflePartitioner>>.

I have also done some implementation work for Uniffle RSS, and with that approach the native code compiles successfully:

https://github.com/wForget/arrow-datafusion-comet/blob/ff68265277d12c5600b450a4cc67073791062490/native/shuffle/src/shuffle_writer.rs#L212-L272