jayzhan211 commented on code in PR #26960:
URL: https://github.com/apache/datafusion/pull/26960#discussion_r4454691795
##########
datafusion/execution/src/disk_manager.rs:
##########
@@ -578,7 +584,7 @@ impl SpillFile for RefCountedTempFile {
.try_clone()
.map_err(DataFusionError::IoError)?;
Ok(Box::new(FileSpillWriter {
- file,
+ file: BufWriter::with_capacity(SPILL_WRITE_BUFFER_SIZE, file),
Review 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` at 8192-row batches.
```diff
--- 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`.)
##########
datafusion/physical-plan/src/spill/mod.rs:
##########
@@ -677,14 +700,25 @@ impl IPCStreamEncoder {
let compression_type =
Option::<CompressionType>::from(spill_compression);
write_options = write_options.try_with_compression(compression_type)?;
- let encoder = StreamEncoder::try_new_with_options(schema,
write_options)?;
+ // Lengths take as much space as offsets when uncompressed
+ let lengths_schema =
+ compression_type.and_then(|_| lengths::encode_schema(schema));
Review 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. At 1–8 rows per batch, files roughly double. 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?
```diff
+/// 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:
```rs
#[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(())
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
