Repository navigation
Conversation
The existing multi-input test never spilled, so nothing checked that each output partition's range stitches every input's spill and in-memory bytes together in order. Run it again with a tiny buffer budget so every input spills several times and keeps an in-memory remainder.
The spill manager kept one IPC writer open per spilled output partition until the task finished, each holding a file descriptor and its own compression context outside the memory pool. With hundreds of output partitions that is hundreds of open files per input partition and, with zstd, hundreds of megabytes of compressor state. Write each spill event as one complete stream appended to the partition's file and close it; readers already cross concatenated streams.
|
@NoahKusaba could you help review, since you are familiar with shuffle? |
There was a problem hiding this comment.
I've read through everything, and it's a clean refactor. I did notice some opportunities for cleanups, which I added as non-blocking comments.
I did spend some time exploring an alternative design where each spill writes one file, rather than one file per output partition (up to M x K files on disk). That means far fewer files and file opens, at the cost of more, smaller copies at finalize and a small in-memory table of partition lengths, so performance should roughly break even; the win is simplicity.
I think we are all working on higher leverage changes regardless, so it's just a fun thought 😄 .
I did have opus look into it, and here's what it had to say (This is how Spark does it):
Spark's ShuffleExternalSorter instead writes one file per spill, with partitions in order, and keeps only their lengths (SpillInfo.partitionLengths). mergeSpillsWithTransferTo then builds each output partition by copying its slice out of every spill file in turn. That works because compressed streams can be concatenated, as our IPC streams already are.
For us that would mean:
- Spill: write each partition as its own IPC stream into one file per spill, and record
Vec<u64>lengths. - Finalize: for each partition,
copy_file_rangeits slice from each spill file. The readers and the output format don't change. - Code: spilling and in-memory encoding become one helper writing to two places, and
SpillManagershrinks to a list of(file, lengths).
Cost: I'd expect performance to roughly break even. The bytes copied are the same; we trade per-spill opens for more, smaller copies at finalize. Spark's own docs note that those small copies hurt when partitions per spill are small. So I'd check a spill-heavy query with a small memory_limit_per_task_bytes before switching.
| partition_id: usize, | ||
| ) -> Result<Vec<RecordBatch>> { | ||
| use datafusion::arrow::ipc::reader::StreamReader; | ||
| use std::io::{Seek, SeekFrom}; |
There was a problem hiding this comment.
nit - I've seen in other repo's the precedence of keeping imports at the file header and not in-lining them in functions.
There was a problem hiding this comment.
Done in 7d51a94. The inline imports moved up to the test module header.
| return Ok(0); | ||
| /// Returns `(batches, bytes)` written, with bytes estimated from each | ||
| /// batch's array memory size. | ||
| pub fn spill_batches( |
There was a problem hiding this comment.
This loop is now essentially encode_buffered_partitions (writer.rs:121): write one partition's batches as a single IPC stream and count batches, rows and bytes. The only difference is the sink: a file opened for appending here, a Vec<u8> there. Could both call one shared helper?
fn write_stream<W: Write>(
w: W,
schema: &SchemaRef,
opts: &IpcWriteOptions,
batches: &mut dyn Iterator<Item = Result<RecordBatch>>,
) -> Result<(W, EncodedStats)>Spilling would then be: open the file in append mode, call write_stream with a BufWriter, flush, and add the stats to the partition's totals. partition_counters could hold EncodedStats instead of a (u64, u64, u64) tuple, since it stores the same three counts.
Also, SpillManager could build its IpcWriteOptions once in new() instead of calling try_with_compression on every spill (line 144).
There was a problem hiding this comment.
Done in 7d51a94. Spill and the in-memory encode now share write_ipc_stream, partition_counters holds EncodedStats, and the options are built once in new().
| batches: &mut dyn Iterator<Item = datafusion::error::Result<RecordBatch>>, | ||
| ) -> Result<(u64, u64)> { | ||
| let spill_path = self.spill_file_path(partition_id); | ||
| let mut writer: Option<StreamWriter<BufWriter<std::fs::File>>> = None; |
There was a problem hiding this comment.
nit - I don't think the lazy open or the num_rows() == 0 skip can be reached from spill_all_partitions. BufferedBatches::drain only calls f for partitions with at least one row, and PartitionedBatchIterator never returns an empty batch. The only caller that could pass an empty batch is the spill() wrapper.
If that wrapper returns early on an empty batch, the writer can be created before the loop, and the Option and match go away. Both functions get shorter.
There was a problem hiding this comment.
Done in 7d51a94. The writer is created up front now, so the Option and the zero-row skip are gone.
| /// Schema shared by all spill writers | ||
| schema: SchemaRef, | ||
| /// Spill file path per output partition: partition_id -> spill_file_path | ||
| spill_files: HashMap<usize, PathBuf>, |
There was a problem hiding this comment.
spill_files and partition_counters now track the same thing: each gets an entry exactly when a partition has spilled at least once. The stored path is also just spill_file_path(partition_id). And now that this runs per spill, it rebuilds the PathBuf on every call (line 125) and re-inserts it (line 171), not only on the first spill.
Could we drop spill_files? has_spill_files could check partition_counters.contains_key, and spill_path could compute the path, returning PathBuf, or Option<PathBuf> when the partition has spilled.
There was a problem hiding this comment.
Done in 7d51a94. spill_files is gone. has_spill_files checks partition_counters and spill_path computes an Option<PathBuf>.
| /// | ||
| /// Returns the number of bytes written (estimated from the batch's array | ||
| /// memory size). | ||
| pub fn spill(&mut self, partition_id: usize, batch: &RecordBatch) -> Result<u64> { |
There was a problem hiding this comment.
nit - Only the tests call spill now. Since this PR already has api-change, how about removing it and having the tests call spill_batches(id, &mut std::iter::once(Ok(batch)))? The alternative is #[cfg(test)]. Either way, the public API doesn't keep a single-batch wrapper that clones its input on every call.
There was a problem hiding this comment.
Done in 7d51a94. spill() is removed and the tests call spill_batches through a small helper.
Address review feedback: spill and the in-memory encode now call one write_ipc_stream helper, SpillManager builds its IpcWriteOptions once, derives spill paths from partition_counters instead of a duplicate map, and the single-batch spill() wrapper is removed (tests call spill_batches).
|
@NoahKusaba Thanks for digging into the Spark layout! I'd like to keep this PR scoped to closing the writers, so I opened #2810 to look into it separately, including the spill-heavy benchmark at a small |
|
lgtm, approving! |
|
Thanks for approving @NoahKusaba. Since these changes are now also in #2805 which has committer approval, I am merging that one and closing this one. |
Which issue does this PR close?
No issue. Split out of #2792. It is a prerequisite for making zstd the default shuffle codec.
Rationale for this change
The sort shuffle spill manager keeps one Arrow IPC
StreamWriteropen for every output partition it has spilled, until the input partition finishes. Each open writer holds a file descriptor and its own compression context, and the contexts are not tracked by the memory pool. With LZ4 the contexts are small. With zstd each one is about 0.3–0.7 MB. Once an input spills, it usually spills every output partition, so a task with M local inputs and K output partitions holds up to M × K writers at once. At 16 inputs and 256 partitions that is 4,096 open files per task and about 2.9 GB of untracked zstd contexts per executor.What changes are included in this PR?
SpillManager::spill_batcheswrites one spill event's batches for a partition as one complete IPC stream, appended to that partition's spill file, and closes the writer before returning. A spill file therefore holds one stream per spill event. The consolidated output already concatenates several streams per partition range, and the readers (MultiStreamPartitionStreamand the remote fetch path) already read across them, so the output file format is unchanged.spillstays, as a one-batch wrapper.finish_writersis removed, since no writer outlives a call.open_spill_readeris removed too: it read only the first stream, and nothing outside the tests used it.multiple_spilling_input_partitions_round_trip: several input partitions each spill repeatedly and keep an in-memory remainder, and every row must come back exactly once with its value. The existing multi-input test never spilled, so nothing checked how each output partition's range stitches spill and in-memory bytes from every input together.On TPC-H SF100 (2 executors × 8 cores, 16 partitions, zstd) the suite total is unchanged: 132.8 s against 132.3 s.
Are there any user-facing changes?
No behavior changes.
SpillManageris public, andfinish_writersandopen_spill_readerare removed from it, so this has theapi-changelabel.