feat: add destination-aware native shuffle execution (3/n)

native/shuffle/src/lib.rs

@@ -19,6 +19,8 @@ pub(crate) mod comet_partitioning;
 pub mod ipc;
 pub(crate) mod metrics;
 pub(crate) mod partitioners;
+#[cfg(test)]
+mod rss_execution_tests;
 mod schema_align;
 mod shuffle_writer;
 mod spark_crc32c_hasher;
@@ -28,5 +30,5 @@ pub(crate) mod writers;
 pub use comet_partitioning::CometPartitioning;
 pub use ipc::read_ipc_compressed;
 pub use schema_align::SchemaAlignExec;
-pub use shuffle_writer::ShuffleWriterExec;
+pub use shuffle_writer::{ShuffleWriterDestination, ShuffleWriterExec};
 pub use writers::{CompressionCodec, ShuffleBlockWriter};

native/shuffle/src/metrics.rs

@@ -33,7 +33,8 @@ pub(crate) struct ShufflePartitionerMetrics {
     /// Time encoding batches to IPC format
     pub(crate) encode_time: Time,
 
-    /// Time spent writing to disk. Maps to "shuffleWriteTime" in Spark SQL Metrics.
+    /// Time spent writing encoded data to its destination. Maps to "shuffleWriteTime" in Spark
+    /// SQL Metrics.
     pub(crate) write_time: Time,
 
     /// Number of input batches

native/shuffle/src/partitioners/empty_schema.rs

@@ -26,7 +26,7 @@ use tokio::time::Instant;
 /// A partitioner for zero-column schemas (e.g. queries where ColumnPruning removes all columns).
 /// This handles shuffles for operations like COUNT(*) that produce empty-schema record batches
 /// but contain a valid row count. Accumulates the total row count and writes a single
-/// zero-column IPC batch to partition 0. All other partitions get empty entries in the index file.
+/// zero-column IPC batch to partition 0. All other partitions are finalized without data.
 pub(crate) struct EmptySchemaShufflePartitioner<T: PartitionWriter> {
     partition_writer: T,
     schema: SchemaRef,

native/shuffle/src/partitioners/multi_partition.rs

@@ -554,7 +554,7 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
 
     pub(crate) fn spill(&mut self, unreserved_bytes: usize) -> datafusion::common::Result<()> {
         log::info!(
-            "ShuffleRepartitioner spilling shuffle data of {} to disk while inserting ({} time(s) so far)",
+            "ShuffleRepartitioner spilling {} bytes to its partition writer ({} previous spills)",
             self.used(),
             self.spill_count()
         );

native/shuffle/src/partitioners/single_partition.rs

@@ -22,19 +22,19 @@ use arrow::array::RecordBatch;
 use std::iter;
 use tokio::time::Instant;
 
-/// A partitioner that writes all shuffle data to a single file and a single index file.
+/// A partitioner that streams all shuffle data to output partition zero.
 ///
-/// Batches are streamed straight to the long-lived `BufBatchWriter` inside the
-/// [`PartitionWriter`], whose internal `BatchCoalescer` combines sub-`batch_size` batches
-/// into `batch_size`-row IPC blocks across calls. A batch that is already `>= batch_size`
-/// and lands on an empty coalescer buffer is passed through and written verbatim as a
-/// single block, which may exceed `batch_size` (see the `BatchCoalescer` bypass in
-/// `BufBatchWriter`). Block boundaries therefore depend on how the input is chunked, but
-/// every row is written exactly once in order.
+/// Batches are streamed straight to the underlying [`PartitionWriter`]. For local output, its
+/// long-lived `BufBatchWriter` uses an internal `BatchCoalescer` to combine sub-`batch_size`
+/// batches into `batch_size`-row IPC blocks across calls. A batch that is already `>= batch_size`
+/// and lands on an empty coalescer buffer is passed through and written verbatim as a single
+/// block, which may exceed `batch_size` (see the `BatchCoalescer` bypass in `BufBatchWriter`).
+/// Remote output instead sends each encoded batch directly to its task-owned callback. In either
+/// case every row is written exactly once in order.
 ///
-/// The partitioner does no buffering or concatenation of its own. A concat layer here
-/// would be redundant with the coalescer, and actively wasteful whenever its output landed
-/// below `batch_size` and so missed the bypass: those rows would be copied once by the
+/// The partitioner does no buffering or concatenation of its own. For local output, a concat
+/// layer here would be redundant with the coalescer, and actively wasteful whenever its output
+/// landed below `batch_size` and so missed the bypass: those rows would be copied once by the
 /// concat and again into the coalescer's builders. That happens whenever the input batch
 /// size does not divide `batch_size` evenly. When it does divide evenly the concat output
 /// hit the bypass and only one copy was made either way, so streaming costs nothing there
@@ -64,7 +64,7 @@ impl<T: PartitionWriter> ShufflePartitioner for SinglePartitionShufflePartitione
             self.metrics.data_size.add(batch.get_array_memory_size());
             self.metrics.baseline.record_output(num_rows);
 
-            // Stream directly to the writer; its BatchCoalescer handles batching to batch_size.
+            // Stream directly to the writer; each destination owns its encoding and batching.
             self.partition_writer
                 .write(0, &mut iter::once(Ok(batch)), &self.metrics)?;
         }