Repository navigation
Conversation
`StreamEncoder` returns every buffer of a batch and its padding as a separate `Buffer`, and the spill writers wrote them one by one. For file backed spills that is one `write` system call per buffer, about 35 for a batch of six string and integer columns. Copy runs of buffers under 64 KB into one, so a spilled batch typically takes a single write.
Closes apache#26627. Offsets are a running total, so zstd keeps about 2.5 of every 4 bytes and lz4 all of them, while the lengths of similar values repeat. Compressed spill files now store each top level `Utf8`, `LargeUtf8`, `Binary` and `LargeBinary` column as a struct of its lengths and of its value bytes as a single run. Both are valid arrays, so the IPC reader decodes them as usual and the spill reader rebuilds the offsets with a running sum.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
- `FileSpillWriter` writes through a 256 KB `BufWriter` instead of the IPC encoder copying small buffers together. One bounded buffer is reused, large buffers are not copied, and uncompressed spills benefit too. - Byte arrays read back from compressed spill files are rebuilt with `OffsetBuffer::try_from_lengths` and `GenericByteArray::try_new`, so a corrupted file fails instead of producing invalid arrays. - `IPCStreamEncoder` keeps the encoded schema and reuses its struct fields for every batch. - `aggregate_memory_spill.slt` spills byte array keys with zstd and lz4.
github-actions
Bot
added
sqllogictest
comphead
changed the title
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
@comphead thanks for working on this, the size wins at normal batch sizes are great. Two things to address before merging, details inline: the per-file write buffer is memory the pool can't see, and small batches now spill larger than on main.
| .map_err(DataFusionError::IoError)?; | ||
| Ok(Box::new(FileSpillWriter { | ||
| file, | ||
| file: BufWriter::with_capacity(SPILL_WRITE_BUFFER_SIZE, file), |
There was a problem hiding this comment.
@comphead thanks for this! with_capacity allocates the 256 KB up front, and every open spill writer keeps it until it finishes, without the memory pool knowing. A preserve-order RepartitionExec keeps one spill file open per (input, output) pair, so this grows quadratically with partitions. With 2 inputs × 128 outputs and a 512 KB pool, peak heap went from 2.35 MB on main to 69.5 MB with this PR, and aggregate_memory_spill.slt now reports about twice the peak-memory drift. Your first commit's per-batch coalescing avoids this, since its copy only lives for one write. I ran it in place of the BufWriter: heap is back to main's level, the spill tests pass, and writes match or beat the BufWriter in most 8192-row cases.
--- a/datafusion/execution/src/disk_manager.rs
+++ b/datafusion/execution/src/disk_manager.rs
- file: BufWriter::with_capacity(SPILL_WRITE_BUFFER_SIZE, file),
+ file,
--- a/datafusion/physical-plan/src/spill/mod.rs
+++ b/datafusion/physical-plan/src/spill/mod.rs
- Ok((batch.num_rows(), bytes, buffers))
+ Ok((batch.num_rows(), bytes, coalesce_buffers(buffers)))
...
- Ok((bytes, buffers))
+ Ok((bytes, coalesce_buffers(buffers)))
+
+/// Buffers at least this large are written on their own, see [`coalesce_buffers`]
+const COALESCE_BUFFERS_BELOW: usize = 64 * 1024;
+
+/// Copies runs of small buffers into one, so that writing a batch takes a few
+/// writes instead of one per IPC buffer and padding. The copy lives only for
+/// the write, unlike a buffer kept by every open spill file.
+fn coalesce_buffers(buffers: Vec<Buffer>) -> Vec<Buffer> {
+ let mut coalesced = Vec::new();
+ let mut small: Vec<u8> = Vec::new();
+ for buffer in buffers {
+ if buffer.len() < COALESCE_BUFFERS_BELOW {
+ small.extend_from_slice(buffer.as_slice());
+ continue;
+ }
+ if !small.is_empty() {
+ coalesced.push(Buffer::from_vec(std::mem::take(&mut small)));
+ }
+ coalesced.push(buffer);
+ }
+ if !small.is_empty() {
+ coalesced.push(Buffer::from_vec(small));
+ }
+ coalesced
+}(Also revert the finish() flush and the BufWriter import in disk_manager.rs.)
There was a problem hiding this comment.
Thanks, agreed. Restored the per-batch coalescing from your diff in 03a6c93, so disk_manager.rs is unchanged now. aggregate_memory_spill.slt reports a 2.2 MB peak-memory drift with it, against 5.6 MB with the BufWriter.
| // Lengths take as much space as offsets when uncompressed | ||
| let lengths_schema = | ||
| compression_type.and_then(|_| lengths::encode_schema(schema)); |
There was a problem hiding this comment.
The struct and run-end wrapper costs about 200 bytes per batch for each byte column, so small batches now spill larger than on main. A preserve-order hash repartition into 128 partitions (about 64 rows per batch) spills 13.5% more with zstd. Files roughly double at 1 row per batch and grow 44–77% at 8 rows. Gating on the file's first batch fixed it in my run: 128 partitions spill the same as main, and the 64-partition case still wins (lz4 14.0 → 12.9 bytes/row). Could you add this along with the boundary test?
+/// Fewer rows per batch than this and the struct and run-end metadata of a
+/// lengths column outweighs what the lengths save, see [`lengths`]
+const MIN_LENGTHS_ROWS: usize = 128;
...
- fn new(schema: &Schema, spill_compression: SpillCompression) -> Result<Self> {
+ fn new(
+ schema: &Schema,
+ spill_compression: SpillCompression,
+ first_batch_rows: usize,
+ ) -> Result<Self> {
...
- let lengths_schema =
- compression_type.and_then(|_| lengths::encode_schema(schema));
+ let lengths_schema = compression_type
+ .filter(|_| first_batch_rows >= MIN_LENGTHS_ROWS)
+ .and_then(|_| lengths::encode_schema(schema));Pass gc_batch.num_rows() from both append_batch and append_batch_async in in_progress_spill_file.rs, and usize::MAX from the tests. Boundary test, so the encoding is off at 127 rows and on at 128:
#[test]
fn test_spill_lengths_skipped_for_small_batches() -> Result<()> {
let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)]));
for (rows, encoded) in [(127, false), (128, true)] {
let values = (0..rows).map(|i| format!("item-{}", i / 64));
let batch = RecordBatch::try_new(
Arc::clone(&schema),
vec![Arc::new(StringArray::from_iter_values(values))],
)?;
let spill_manager = spill_manager(Arc::clone(&schema), SpillCompression::Zstd);
let spill_file = spill_manager
.spill_record_batch_and_finish(std::slice::from_ref(&batch), "Test")?
.unwrap();
let file = std::fs::File::open(spill_file.path().unwrap())?;
let reader = arrow_ipc::reader::StreamReader::try_new(file, None)?;
let spilled_type = reader.schema().field(0).data_type().clone();
assert_eq!(spilled_type != DataType::Utf8, encoded, "{rows} rows");
}
Ok(())
}There was a problem hiding this comment.
Added MIN_LENGTHS_ROWS and the boundary test as you wrote them, in 03a6c93. The #26627 reproducer with batch_size = 64 now spills the same bytes as main with both codecs, and at 128 rows it still saves 25% with lz4 (70.3 to 52.5 MB) and 1% with zstd. I also set batch_size = 8192 for the byte key case in aggregate_memory_spill.slt, since its 128-row batches sat exactly at the threshold.
- Copy the small IPC buffers of each batch into one write again, instead of keeping a 256 KB `BufWriter` per spill file. Those buffers were outside the memory pool, and a preserve-order `RepartitionExec` keeps one spill file open per input and output partition. - Write lengths only when the first batch of a file has at least 128 rows. Below that, the struct and run-end metadata cost more than the lengths save.
comphead
changed the title
Which issue does this PR close?
Utf8offsets, which zstd and lz4 barely compress #26627.Rationale for this change
Compressed spill files are mostly string offsets, which zstd and lz4 barely compress, see #26627. Spilling is also slower than it needs to be: a batch of the table in #26627 takes about 35
writecalls, one per IPC buffer.This replaces #26698, which needed spill reads to skip validation (#26688).
What changes are included in this PR?
lz4_frameorzstd, each top levelUtf8,LargeUtf8,BinaryandLargeBinarycolumn is spilled as a struct of its value lengths and of its value bytes as one run-end encoded value. Both are valid arrays, so the IPC reader decodes them as usual before the spill reader rebuilds and validates the original arrays. Files whose first batch has fewer than 128 rows keep their offsets, since the struct and run-end metadata would cost more than the lengths save. Uncompressed spills keep their layout.Benchmarks
Apple M5 Max, interleaved runs against
main.The reproducer from #26627 (
datafusion-cli -m 100m, median of 7):The same query with small batches (
batch_size64, 128 and 256, median of 3), spilled bytes and time:sort_tpchqueries at SF1 onUtf8columns (one partition,-m 200m, median of 5), spilled bytes and time againstmain:In the
spill_iobenchmark, writing and reading spill files is 2-14% faster uncompressed, 4-8% faster with zstd and unchanged to 10% faster with lz4, and the string datasets spill 27-41% smaller.spill_ioresults and setupcargo bench -p datafusion-physical-plan --bench spill_io, median of three alternating runs per build.The
sort_tpchruns set:What is the testing strategy for this PR?
aggregate_memory_spill.sltspillsUtf8,LargeUtf8,BinaryandLargeBinarygroup keys with nulls underzstdandlz4_frameand checks their values.test_spill_lengths_round_tripcovers sliced and empty batches,test_spill_lengths_compresschecks that a zstd spill of similar strings is under a quarter of plain zstd IPC, andtest_spill_lengths_skipped_for_small_batcheschecks the 128-row threshold.test_async_ipc_stream_writerchecks that each batch takes one write.Are there any user-facing changes?
Smaller compressed spill files and faster spilling. No configuration or API changes.