diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 13bc2f3def1..63eabc1c30a 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -582,14 +582,21 @@ diverge for several structural reasons: retained/dirty page cache all add resident bytes that no layer above the allocator can see. Freeing memory does not necessarily return pages to the OS. - **Non-Rust allocations.** Memory allocated by C dependencies through libc `malloc`, and anything - `mmap`ed, never passes through Rust's `GlobalAlloc`, so neither the memory pool nor the - allocation counters (`native_allocated`, `jemalloc_allocated`) see it. In a default build the C - dependencies are libzstd (`zstd-sys`, behind the Parquet `zstd` codec), libhdfs (`hdfs-sys`, - pulled in by the default `hdfs-opendal` feature), and the TLS stack used for cloud object stores - (`aws-lc-sys`). With the `python-udf` feature, the embedded Python interpreter and PyArrow also - allocate outside `GlobalAlloc`; their memory is absent from the executor's native `allocated` - figure, and Python working allocations are not reserved in Comet's pool. Building with the - `jemalloc` or `mimalloc` feature adds the allocator itself (`tikv-jemalloc-sys`, + `mmap`ed, never passes through Rust's `GlobalAlloc`, so the allocation counters + (`native_allocated`, `jemalloc_allocated`) never see it, and the memory pool tracks it only where + an operator charges it by hand. The one place that happens today is the native shuffle writer's + zstd context when a task writes to more than one local partition: the writer reserves the + context's size for as long as it holds one, so those bytes are in `reserved` but not in + `allocated`. It holds one between blocks only up to level 8. A larger context exceeds the 8 MiB + the writer keeps, so `release_zstd_if_oversized` frees it after each block, and that workspace, + about 15 MiB at level 9 and 90 MiB at level 19, is in neither figure. In a default build the C + dependencies are libzstd (`zstd-sys`, behind the Parquet + `zstd` codec and the shuffle writer's `zstd` codec), libhdfs (`hdfs-sys`, pulled in by the default + `hdfs-opendal` feature), and the TLS stack used for cloud object stores (`aws-lc-sys`). With the + `python-udf` feature, the embedded Python interpreter and PyArrow also allocate outside + `GlobalAlloc`; their memory is absent from the executor's native `allocated` figure, and Python + working allocations are not reserved in Comet's pool. Building with the `jemalloc` or `mimalloc` + feature adds the allocator itself (`tikv-jemalloc-sys`, `libmimalloc-sys`). It is worth knowing which dependencies are _not_ C, because several names suggest otherwise: the other Parquet codecs are pure Rust in this build, `snap` for Snappy, `lz4_flex` for LZ4 and `zlib-rs` for gzip, as is `libbz2-rs-sys` despite its name, so those diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index a9d47f0f572..b605fcab4ac 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -483,6 +483,26 @@ The `MultiPartitionShuffleRepartitioner` holds: - `max_buffer_bytes`, the optional fixed spill threshold described above. `None` leaves pool pressure as the only trigger. - `scratch`, reusable buffers for partition ID computation. +- Through the writer, a second reservation for the local writer's own buffers: the data file + and spill file write buffers, the spill copy buffer, the encode scratch and the cached zstd + context. `try_new` hands `LocalPartitionWriter` a `reservation.new_empty()` sibling, so the + bytes land on the same consumer without a new one, and the writer keeps it equal to what + those buffers hold. A spill frees `reservation` but not this one, and it counts toward + neither `memory_spilled_bytes` nor `spark.comet.shuffle.native.maxBufferBytes`. Each file + and copy buffer is allocated when the phase that uses it starts: the spill file buffer when + the first spill opens the file, the other two at the first `finish_partition`. That call + flushes the spill file and releases its buffer, then sizes the copy buffer, which serves + every spilled range, and then the data file buffer, which blocks of the write buffer size + bypass anyway. Each asks the pool for its full size, and when the pool refuses the writer + uses an 8 KiB buffer instead and charges that. Any other pool error, such as an exception + Spark threw while spilling another consumer, fails the task. A spill forced by pool pressure + opens the spill file while its batches are still charged, so that buffer often starts + small. After each spill has freed its batches, the writer drops the drained encode scratch, + which one wide row can leave larger than the pool, and grows the spill file buffer toward + the full size, but to no more than half of what the spill freed, so the next batches have + room. It never shrinks that buffer, and moves the buffered bytes into the larger one. The + scratch and the zstd context already exist when they are charged, so the writer records + them without asking. The spill file is owned by `PartitionedSpill` in `writers/local/spill.rs`. It holds one DataFusion `SpillFile`, created lazily on the first spill and shared by every output partition, and tracks per diff --git a/docs/source/user-guide/latest/tuning/memory.md b/docs/source/user-guide/latest/tuning/memory.md index 2dd730e9365..02b57cb0517 100644 --- a/docs/source/user-guide/latest/tuning/memory.md +++ b/docs/source/user-guide/latest/tuning/memory.md @@ -44,10 +44,16 @@ needs. A native Iceberg fanout write is the exception: it closes some of its ope Comet's memory pool only tracks memory that an operator explicitly reserves, which in practice means the batches an operator deliberately accumulates: the sort buffer, the build side of a hash join, hash aggregation state, the -shuffle writer's buffered partitions, and the files a native writer holds open. Memory that is not reserved is -invisible to the pool no matter how much of it there is. That includes: +shuffle writer's buffered partitions, the files a native writer holds open, and the write buffers, encode scratch +and zstd context of the native shuffle writer when it writes to more than one local partition (the scratch is +charged as each write ends, so its transient peak within a write, the write buffer size plus one block, is not, and +the zstd context only up to level 8, as [Sizing the Overhead from the Memory Usage Log] describes). +Memory that is not reserved is invisible to the pool no matter how much of it there is. That includes: - per-batch working memory in expression kernels and Arrow array builders, +- the write buffers and encode scratch of the single-partition, empty-schema and remote shuffle writers + (`spark.comet.shuffle.native.writeBufferSize` sizes the local ones), and the zstd context they keep while + encoding when `spark.comet.shuffle.compression.codec` is `zstd`, - decompression buffers and Parquet reader structures, - object store request buffers and the async runtime's own machinery, - Arrow buffers allocated on the JVM side, which no budget covers at all, @@ -182,11 +188,19 @@ alongside the JVM's own non-heap memory. To size the overhead from it: non-heap memory, and add the most untracked memory seen on any executor. 3. Add a margin on top. The log can miss the true peak between samples, and none of the figures includes the allocator's fragmentation and retained pages, or memory allocated by native C - libraries such as zstd. With native Arrow UDFs, budget the embedded Python interpreter and - PyArrow allocations in addition to the untracked memory calculated from the log. A PyArrow - result buffer later imported into the JVM can appear in the `imported from native` figure that - the estimate subtracts, even though Rust's `allocated` does not count it. Use the executor's - peak resident memory to size this additional overhead. + libraries such as zstd. The zstd context of the native shuffle writer is the exception at levels + up to 8: when `spark.comet.shuffle.compression.codec` is `zstd`, a shuffle task writing to more + than one local partition reserves the context it keeps between blocks while spilling or writing + its output, so that memory is in `reserved` but never in `allocated`, and the untracked memory + worked out from the log understates what has to fit in the overhead by one context per such task + (about 1.3 MiB at the default level, close to 8 MiB at levels 7 and 8). Above level 8 the context + is too large to keep, so each block creates and frees its own and it appears in neither figure: + budget one per task writing shuffle output, about 15 MiB at level 9 and 90 MiB at level 19, on + top of the untracked memory. With native Arrow UDFs, budget the embedded + Python interpreter and PyArrow allocations in addition to the untracked memory calculated from + the log. A PyArrow result buffer later imported into the JVM can appear in the + `imported from native` figure that the estimate subtracts, even though Rust's `allocated` does + not count it. Use the executor's peak resident memory to size this additional overhead. For example, a 16 GiB executor derives an overhead of 1638 MiB. If the line above has the most untracked memory in its log, that is 5412.3 + (310.4 - 96.2) - 3890.0 = 1736.5 MiB. The overhead diff --git a/native/shuffle/src/codec_context.rs b/native/shuffle/src/codec_context.rs index 114f0530252..dd15b490cad 100644 --- a/native/shuffle/src/codec_context.rs +++ b/native/shuffle/src/codec_context.rs @@ -112,6 +112,10 @@ impl ShuffleCodecContext { } } + pub(crate) fn retained_bytes(&self) -> usize { + self.zstd.as_ref().map_or(0, |cctx| cctx.sizeof()) + } + /// Test hook for the release-vs-retain contract of the two encode paths. #[cfg(test)] pub(crate) fn holds_zstd_cctx(&self) -> bool { diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index f217165d38f..b73a7fc7a83 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -180,7 +180,7 @@ impl MultiPartitionShuffleRepartitioner { #[allow(clippy::too_many_arguments)] pub(crate) fn try_new( partition: usize, - partition_writer: T, + mut partition_writer: T, partitioning: CometPartitioning, metrics: ShufflePartitionerMetrics, runtime: Arc, @@ -235,6 +235,9 @@ impl MultiPartitionShuffleRepartitioner { let reservation = MemoryConsumer::new(format!("ShuffleRepartitioner[{partition}]")) .with_can_spill(true) .register(&runtime.memory_pool); + // The writer's buffers share the consumer but not the reservation: `spill` frees this + // one and reports it as spilled memory, and the buffer limit is checked against it. + partition_writer.attach_buffer_reservation(reservation.new_empty())?; Ok(Self { buffered_batches: vec![], @@ -644,17 +647,18 @@ impl MultiPartitionShuffleRepartitioner { ) }) }; - self.partition_writer.write_burst_complete(); - // Count the input capacity released from buffering by this spill, including a // rejected reservation. Shared allocations are charged once within a spill, but // contribute again if buffered for a later spill, regardless of input batching. // Also release and count all buffered inputs when the writer fails partway through. - let memory_spilled_bytes = self.reservation.free().saturating_add(unreserved_bytes); + let freed = self.reservation.free(); + let memory_spilled_bytes = freed.saturating_add(unreserved_bytes); self.metrics.memory_spilled_bytes.add(memory_spilled_bytes); self.pinned_buffers.clear(); + // After the free, so the writer can regrow buffers the pool cut short during the spill. + let burst_result = self.partition_writer.write_burst_complete(freed); self.metrics.spill_count.add(1); - write_result + write_result.and(burst_result) }) } @@ -734,13 +738,18 @@ impl Debug for MultiPartitionShuffleRepartitioner { #[cfg(test)] mod tests { use super::*; + use crate::writers::LocalPartitionWriter; + use crate::{CompressionCodec, PartitionOffsets, ShuffleBlockWriter}; use arrow::array::{ArrayRef, Int64Array}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; + use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; #[derive(Default)] struct FailingPartitionWriter { fail: bool, consume_before_failure: bool, + fail_burst: bool, write_calls: usize, } @@ -784,6 +793,59 @@ mod tests { ) -> datafusion::common::Result<()> { Ok(()) } + + fn write_burst_complete(&mut self, _freed: usize) -> datafusion::common::Result<()> { + if self.fail_burst { + return Err(DataFusionError::Execution( + "injected burst failure".to_string(), + )); + } + Ok(()) + } + } + + /// A failed `write_burst_complete` fails the spill, but only after the spilled memory is + /// released and counted; a write error takes precedence. + #[tokio::test] + async fn spill_returns_burst_error_after_releasing_memory() { + let batch = int64_batch(0..64); + let runtime = Arc::new(RuntimeEnv::default()); + let mut repartitioner = hash_repartitioner( + FailingPartitionWriter { + fail_burst: true, + ..Default::default() + }, + 2, + Arc::clone(&runtime), + None, + ); + for (fail_write, expected) in [ + (false, "injected burst failure"), + (true, "injected write failure"), + ] { + repartitioner.partition_writer.fail = fail_write; + repartitioner.insert_batch(batch.clone()).await.unwrap(); + let reserved = repartitioner.reservation.size(); + assert!(reserved > 0); + let (spills, spilled) = ( + repartitioner.spill_count(), + repartitioner.metrics.memory_spilled_bytes.value(), + ); + + let error = repartitioner.spill(0).unwrap_err(); + assert!( + matches!(&error, DataFusionError::Execution(message) if message == expected), + "unexpected error: {error}" + ); + assert_eq!(runtime.memory_pool.reserved(), 0); + assert_eq!( + ( + repartitioner.spill_count(), + repartitioner.metrics.memory_spilled_bytes.value() + ), + (spills + 1, spilled + reserved) + ); + } } async fn check_spill_metrics_count_input_buffers(batch: RecordBatch, input_bytes: usize) { @@ -997,8 +1059,6 @@ mod tests { #[tokio::test] async fn rejected_growth_counts_existing_reservation_and_unreserved_input() { - use datafusion::execution::runtime_env::RuntimeEnvBuilder; - for share_buffer in [false, true] { for (fail, consume_before_failure) in [(false, false), (true, false), (true, true)] { let batch = RecordBatch::try_from_iter([( @@ -1063,8 +1123,6 @@ mod tests { async fn heterogeneous_spill_metrics_do_not_depend_on_input_batching() { use arrow::array::{DictionaryArray, Int32Array, ListArray, StringArray, StringViewArray}; use arrow::datatypes::{Int32Type, Int64Type}; - use datafusion::physical_expr::expressions::Column; - let num_rows = 64usize; let batch_size = 4usize; let strings = (0..num_rows) @@ -1355,8 +1413,6 @@ mod tests { /// each partition still in input order. #[tokio::test] async fn positional_placement_survives_spilling() { - use datafusion::execution::runtime_env::RuntimeEnvBuilder; - // Ragged, so that groups and slices straddle batch and spill boundaries. let framing = [300, 17, 683, 64, 1, 191]; let (num_partitions, group_rows, start_partition, batch_size) = (8, 48, 3, 256); @@ -1436,6 +1492,528 @@ mod tests { ); } + fn runtime_with_pool(limit: usize) -> Arc { + Arc::new( + RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .build() + .unwrap(), + ) + } + + fn local_writer( + dir: &tempfile::TempDir, + num_partitions: usize, + write_buffer_size: usize, + runtime: Arc, + ) -> LocalPartitionWriter { + local_writer_for( + int64_batch(0..1).schema().as_ref(), + dir, + num_partitions, + write_buffer_size, + runtime, + ) + } + + fn local_writer_for( + schema: &arrow::datatypes::Schema, + dir: &tempfile::TempDir, + num_partitions: usize, + write_buffer_size: usize, + runtime: Arc, + ) -> LocalPartitionWriter { + let block_writer = ShuffleBlockWriter::try_new(schema, CompressionCodec::Lz4Frame).unwrap(); + LocalPartitionWriter::try_new( + dir.path().join("data.out").to_str().unwrap().to_string(), + Arc::new(PartitionOffsets::default()), + block_writer, + num_partitions, + 1024, + write_buffer_size, + runtime, + ) + .unwrap() + } + + fn hash_repartitioner( + writer: W, + num_partitions: usize, + runtime: Arc, + max_buffer_bytes: Option, + ) -> MultiPartitionShuffleRepartitioner { + try_hash_repartitioner(writer, num_partitions, runtime, max_buffer_bytes).unwrap() + } + + fn try_hash_repartitioner( + writer: W, + num_partitions: usize, + runtime: Arc, + max_buffer_bytes: Option, + ) -> datafusion::common::Result> { + MultiPartitionShuffleRepartitioner::try_new( + 0, + writer, + CometPartitioning::Hash(vec![Arc::new(Column::new("v", 0))], num_partitions), + ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + runtime, + 1024, + false, + max_buffer_bytes, + ) + } + + /// A pool error that is not a refusal, such as an exception Spark threw while spilling a + /// peer, fails the shuffle instead of leaving the writer on a fallback buffer. + #[tokio::test] + async fn buffer_pool_errors_other_than_refusal_fail_the_shuffle() { + use crate::writers::erroring_pool::{external, ErroringPool}; + let write_buffer_size = 1 << 20; + let dir = tempfile::tempdir().unwrap(); + // Only the writer's buffers ask for the full write buffer size. + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(ErroringPool::new( + 1 << 30, + write_buffer_size, + external, + ))) + .build() + .unwrap(), + ); + let result = async { + let mut repartitioner = try_hash_repartitioner( + local_writer(&dir, 2, write_buffer_size, Arc::clone(&runtime)), + 2, + Arc::clone(&runtime), + None, + )?; + repartitioner.insert_batch(int64_batch(0..4096)).await?; + repartitioner.spill(0)?; + repartitioner.shuffle_write() + } + .await; + let error = result.expect_err("an exception from the pool must fail the shuffle"); + assert!( + matches!(error.find_root(), DataFusionError::External(_)), + "unexpected error: {error}" + ); + assert_eq!(runtime.memory_pool.reserved(), 0); + } + + fn assert_local_buffers_charged( + repartitioner: &MultiPartitionShuffleRepartitioner, + runtime: &RuntimeEnv, + ) -> usize { + let writer = repartitioner.partition_writer(); + let charged = writer.buffer_reservation_size(); + assert_eq!( + charged, + writer.buffer_bytes_held(), + "reservation must equal the bytes held" + ); + assert_eq!( + runtime.memory_pool.reserved(), + repartitioner.reservation.size() + charged + ); + charged + } + + #[tokio::test] + async fn local_writer_buffers_are_charged_and_released() { + let write_buffer_size = 1 << 20; + let dir = tempfile::tempdir().unwrap(); + let runtime = runtime_with_pool(1 << 30); + let mut repartitioner = hash_repartitioner( + local_writer(&dir, 2, write_buffer_size, Arc::clone(&runtime)), + 2, + Arc::clone(&runtime), + None, + ); + assert_eq!( + assert_local_buffers_charged(&repartitioner, &runtime), + 0, + "nothing is charged before the first row" + ); + + repartitioner + .insert_batch(int64_batch(0..4096)) + .await + .unwrap(); + repartitioner.spill(0).unwrap(); + assert_eq!(repartitioner.reservation.size(), 0); + assert!( + assert_local_buffers_charged(&repartitioner, &runtime) >= write_buffer_size, + "the spill file buffer is charged after a spill" + ); + + repartitioner + .insert_batch(int64_batch(4096..8192)) + .await + .unwrap(); + repartitioner.shuffle_write().unwrap(); + assert!( + assert_local_buffers_charged(&repartitioner, &runtime) >= 2 * write_buffer_size, + "the data file and spill copy buffers are charged after the write" + ); + assert_eq!( + repartitioner.partition_writer().buffer_capacities(), + (write_buffer_size, 0, write_buffer_size), + "the spill file buffer is released before the output is written" + ); + + drop(repartitioner); + assert_eq!(runtime.memory_pool.reserved(), 0); + } + + async fn write_small_input( + dir: &tempfile::TempDir, + limit: usize, + write_buffer_size: usize, + ) -> Vec { + let runtime = runtime_with_pool(limit); + let mut repartitioner = hash_repartitioner( + local_writer(dir, 2, write_buffer_size, Arc::clone(&runtime)), + 2, + Arc::clone(&runtime), + None, + ); + assert_eq!(assert_local_buffers_charged(&repartitioner, &runtime), 0); + // an operator feeding the shuffle, such as a hash join, still gets memory + let upstream = MemoryConsumer::new("HashJoinInput[0]").register(&runtime.memory_pool); + upstream.try_grow(128).unwrap(); + drop(upstream); + repartitioner + .insert_batch(int64_batch(0..16)) + .await + .unwrap(); + repartitioner.shuffle_write().unwrap(); + assert_local_buffers_charged(&repartitioner, &runtime); + std::fs::read(dir.path().join("data.out")).unwrap() + } + + /// The data file buffer is charged only once the output is written, not before the input + /// produces a row. + #[tokio::test] + async fn writer_charges_nothing_before_the_first_row() { + let write_buffer_size = 1 << 20; + let unbounded = + write_small_input(&tempfile::tempdir().unwrap(), 1 << 30, write_buffer_size).await; + let bounded = write_small_input( + &tempfile::tempdir().unwrap(), + write_buffer_size, + write_buffer_size, + ) + .await; + assert!(!unbounded.is_empty()); + assert_eq!(bounded, unbounded); + } + + async fn spill_every_chunk( + dir: &tempfile::TempDir, + runtime: Arc, + write_buffer_size: usize, + ) -> (Vec, (usize, usize, usize)) { + let mut repartitioner = hash_repartitioner( + local_writer(dir, 2, write_buffer_size, Arc::clone(&runtime)), + 2, + Arc::clone(&runtime), + Some(1), + ); + assert_local_buffers_charged(&repartitioner, &runtime); + for start in (0..8192i64).step_by(4096) { + repartitioner + .insert_batch(int64_batch(start..start + 4096)) + .await + .unwrap(); + assert_local_buffers_charged(&repartitioner, &runtime); + } + repartitioner.shuffle_write().unwrap(); + assert_local_buffers_charged(&repartitioner, &runtime); + let capacities = repartitioner.partition_writer().buffer_capacities(); + ( + std::fs::read(dir.path().join("data.out")).unwrap(), + capacities, + ) + } + + #[tokio::test] + async fn refused_buffers_fall_back_to_small_buffers() { + let write_buffer_size = 1 << 20; + let fallback = 8 * 1024; + let dir = tempfile::tempdir().unwrap(); + let (baseline, capacities) = + spill_every_chunk(&dir, runtime_with_pool(1 << 30), write_buffer_size).await; + assert_eq!(capacities, (write_buffer_size, 0, write_buffer_size)); + + for (limit, expected) in [ + // room for one of the output buffers, which goes to the spill copy buffer + ( + write_buffer_size + 512 * 1024, + (fallback, 0, write_buffer_size), + ), + // too small for any of them + (512 * 1024, (fallback, 0, fallback)), + ] { + let dir = tempfile::tempdir().unwrap(); + let (output, capacities) = + spill_every_chunk(&dir, runtime_with_pool(limit), write_buffer_size).await; + assert_eq!(capacities, expected, "pool limit {limit}"); + assert_eq!(output, baseline, "pool limit {limit} changed the output"); + } + } + + fn output_values(bytes: &[u8]) -> Vec { + let mut values = Vec::new(); + let mut pos = 0; + while pos < bytes.len() { + let len = u64::from_le_bytes(bytes[pos..pos + 8].try_into().unwrap()) as usize; + let batch = crate::read_ipc_compressed(&bytes[pos + 16..pos + 8 + len]).unwrap(); + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend_from_slice(column.values()); + pos += 8 + len; + } + values.sort_unstable(); + values + } + + /// 32 MiB of values, so even the 16 MiB pool below spills more than once. + const PRESSURE_ROWS: i64 = 4096 * 1024; + + async fn spill_under_pressure( + dir: &tempfile::TempDir, + runtime: Arc, + write_buffer_size: usize, + ) -> (Vec, Vec) { + let mut repartitioner = hash_repartitioner( + local_writer(dir, 2, write_buffer_size, Arc::clone(&runtime)), + 2, + Arc::clone(&runtime), + None, + ); + let mut spill_buffers = Vec::new(); + for start in (0..PRESSURE_ROWS).step_by(4096) { + let spills = repartitioner.spill_count(); + repartitioner + .insert_batch(int64_batch(start..start + 4096)) + .await + .unwrap(); + if repartitioner.spill_count() > spills { + spill_buffers.push(repartitioner.partition_writer().buffer_capacities().1); + assert_local_buffers_charged(&repartitioner, &runtime); + } + } + repartitioner.shuffle_write().unwrap(); + let bytes = std::fs::read(dir.path().join("data.out")).unwrap(); + (output_values(&bytes), spill_buffers) + } + + /// With no buffer limit every spill is a pressure spill, so the spill file buffer is taken + /// small while the batches still fill the pool. + #[tokio::test] + async fn pressure_spill_restores_the_spill_buffer() { + let write_buffer_size = 256 * 1024; + let expected: Vec = (0..PRESSURE_ROWS).collect(); + for limit in [1 << 20, 4 << 20, 16 << 20] { + let dir = tempfile::tempdir().unwrap(); + let (values, spill_buffers) = + spill_under_pressure(&dir, runtime_with_pool(limit), write_buffer_size).await; + assert!( + spill_buffers.len() >= 2, + "pool limit {limit}: {spill_buffers:?}" + ); + assert!( + spill_buffers.iter().all(|&size| size == write_buffer_size), + "pool limit {limit}: spill buffer after each spill {spill_buffers:?}" + ); + assert!( + values == expected, + "pool limit {limit} lost or repeated rows" + ); + } + } + + /// One row of `v` and `width` incompressible bytes. + fn wide_batch(v: i64, width: usize) -> RecordBatch { + let mut state = 0x9e37_79b9_7f4a_7c15u64 ^ v as u64; + let bytes: Vec = (0..width) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as u8 + }) + .collect(); + RecordBatch::try_from_iter([ + ("v", Arc::new(Int64Array::from(vec![v])) as ArrayRef), + ( + "b", + Arc::new(arrow::array::BinaryArray::from(vec![bytes.as_slice()])) as ArrayRef, + ), + ]) + .unwrap() + } + + /// The scratch that encoded one wide row is released after its spill, so the narrow rows + /// that follow fit the pool instead of each spilling. + #[tokio::test] + async fn wide_row_scratch_is_released_after_the_spill() { + let dir = tempfile::tempdir().unwrap(); + let runtime = runtime_with_pool(512 * 1024); + let mut repartitioner = hash_repartitioner( + local_writer_for( + wide_batch(0, 0).schema().as_ref(), + &dir, + 2, + 1 << 20, + Arc::clone(&runtime), + ), + 2, + Arc::clone(&runtime), + None, + ); + repartitioner + .insert_batch(wide_batch(0, 768 * 1024)) + .await + .unwrap(); + assert_eq!(repartitioner.spill_count(), 1); + assert_local_buffers_charged(&repartitioner, &runtime); + for v in 1..=128 { + repartitioner.insert_batch(wide_batch(v, 0)).await.unwrap(); + } + repartitioner.shuffle_write().unwrap(); + assert!( + repartitioner.spill_count() <= 2, + "{} spills", + repartitioner.spill_count() + ); + let bytes = std::fs::read(dir.path().join("data.out")).unwrap(); + assert_eq!(output_values(&bytes), (0..=128).collect::>()); + } + + /// Feeds 256 batches of 1024 values through 8 partitions, returning the spill count and + /// the output values. After each spill, `check(freed, writer)` sees what the spill freed. + async fn spill_sweep( + writer: W, + runtime: Arc, + mut check: impl FnMut(usize, &W), + ) -> usize { + let mut repartitioner = hash_repartitioner(writer, 8, runtime, None); + for start in (0..256 * 1024i64).step_by(1024) { + let (spills, reserved) = ( + repartitioner.spill_count(), + repartitioner.reservation.size(), + ); + repartitioner + .insert_batch(int64_batch(start..start + 1024)) + .await + .unwrap(); + if repartitioner.spill_count() > spills { + // only a refusal spills here, and a refused batch was never charged + check(reserved, repartitioner.partition_writer()); + } + } + repartitioner.shuffle_write().unwrap(); + repartitioner.spill_count() + } + + /// A restore after a pressure spill grows the spill file buffer to at most half of what the + /// spill freed, leaving the rest to the input, so the writer spills at most about twice as + /// often as one that holds no buffers at all. + #[tokio::test] + async fn spill_buffer_restore_leaves_room_for_the_input() { + for limit in [2112 * 1024, 1088 * 1024] { + let baseline = spill_sweep( + FailingPartitionWriter::default(), + runtime_with_pool(limit), + |_, _| {}, + ) + .await; + let dir = tempfile::tempdir().unwrap(); + let runtime = runtime_with_pool(limit); + let mut previous = 0; + let spills = spill_sweep( + local_writer(&dir, 8, 1 << 20, Arc::clone(&runtime)), + Arc::clone(&runtime), + |freed, writer| { + let capacity = writer.buffer_capacities().1; + // The first spill opens the file, which the pool can grant in full while + // the batches are still charged; after that only a restore grows it. + if previous > 0 && capacity > previous { + assert!( + capacity <= freed / 2, + "pool limit {limit}: spill buffer {capacity} after freeing {freed}" + ); + } + previous = capacity; + assert_eq!(writer.buffer_reservation_size(), writer.buffer_bytes_held()); + }, + ) + .await; + assert!( + spills <= 2 * baseline + 2, + "pool limit {limit}: {spills} spills, {baseline} without buffers" + ); + let bytes = std::fs::read(dir.path().join("data.out")).unwrap(); + assert_eq!( + output_values(&bytes), + (0..256 * 1024i64).collect::>(), + "pool limit {limit}" + ); + } + } + + async fn spill_metrics_with_buffer_limit( + writer: W, + runtime: Arc, + ) -> MultiPartitionShuffleRepartitioner { + let mut repartitioner = hash_repartitioner(writer, 2, runtime, Some(8 * 1024)); + for start in (0..8192i64).step_by(1024) { + repartitioner + .insert_batch(int64_batch(start..start + 1024)) + .await + .unwrap(); + } + repartitioner.shuffle_write().unwrap(); + repartitioner + } + + #[tokio::test] + async fn buffer_charges_do_not_count_as_spilled_memory() { + let baseline = spill_metrics_with_buffer_limit( + FailingPartitionWriter::default(), + runtime_with_pool(1 << 30), + ) + .await; + assert!(baseline.spill_count() > 0); + + let dir = tempfile::tempdir().unwrap(); + let runtime = runtime_with_pool(1 << 30); + let local = spill_metrics_with_buffer_limit( + local_writer(&dir, 2, 1 << 20, Arc::clone(&runtime)), + runtime, + ) + .await; + assert!( + local.partition_writer().buffer_reservation_size() >= 1 << 20, + "the local writer must be charging its buffers for this to mean anything" + ); + assert_eq!( + ( + local.spill_count(), + local.metrics.memory_spilled_bytes.value() + ), + ( + baseline.spill_count(), + baseline.metrics.memory_spilled_bytes.value() + ) + ); + } + /// Positional placement records runs, not rows, and skips the row-level scratch entirely. #[tokio::test] async fn positional_placement_records_runs_not_rows() { diff --git a/native/shuffle/src/writers/local/local_partition_writer.rs b/native/shuffle/src/writers/local/local_partition_writer.rs index d76ff23eff0..1564799f6d9 100644 --- a/native/shuffle/src/writers/local/local_partition_writer.rs +++ b/native/shuffle/src/writers/local/local_partition_writer.rs @@ -23,6 +23,7 @@ use crate::writers::BufBatchWriter; use crate::{PartitionOffsets, ShuffleBlockWriter}; use arrow::array::RecordBatch; use datafusion::common::DataFusionError; +use datafusion::execution::memory_pool::MemoryReservation; use datafusion::execution::runtime_env::RuntimeEnv; use std::fs::{File, OpenOptions}; use std::io::{BufWriter, ErrorKind, Read, Seek, SeekFrom, Write}; @@ -55,19 +56,22 @@ enum DataOutput { /// Multi-partition output: batches are staged per partition and merged into /// `output_writer` one partition at a time during `finish_partition`. Multi { - output_writer: BufWriter, + /// Unbuffered until the first `finish_partition` sizes it, since nothing reaches the + /// data file before then. `None` only while it is rebuilt. + output_writer: Option>, shuffle_block_writer: ShuffleBlockWriter, /// Spilled blocks for every partition, in one file. spill: PartitionedSpill, /// Read handle on the spill file and a write-buffer-sized scratch for its ranges, - /// opened on first use. + /// opened by the first `finish_partition` if anything was spilled. spill_reader: Option<(File, Vec)>, /// Runtime used to allocate the temporary spill file. runtime: Arc, /// Byte buffer recycled through the short-lived per-partition `BufBatchWriter`s. /// Partitions are written strictly one at a time, so a single buffer keeps its - /// grown capacity across the whole task instead of every partition regrowing a - /// fresh allocation toward the write buffer size. + /// grown capacity across the partitions of a spill or of the output instead of every + /// partition regrowing a fresh allocation toward the write buffer size. Each spill + /// releases it. recycled_buffer: Vec, }, } @@ -80,6 +84,9 @@ enum DataOutput { pub(crate) struct LocalPartitionWriter { partition_offsets: Arc, data_output: DataOutput, + /// Kept equal to [`Self::buffer_bytes_held`] at every call boundary. `None` for + /// single-partition output. + buffers: Option, /// Compression state shared by every block this task writes; the per-partition /// `BufBatchWriter`s borrow it (see [`ShuffleCodecContext`]). Retention is bounded: /// released at spill/finish boundaries and whenever its workspace is oversized. @@ -124,7 +131,7 @@ impl LocalPartitionWriter { scratch: Vec::new(), } } else { - let output_writer = BufWriter::with_capacity(write_buffer_size, output_file); + let output_writer = Some(BufWriter::with_capacity(0, output_file)); let spill = PartitionedSpill::new( shuffle_block_writer.clone(), write_buffer_size, @@ -143,6 +150,7 @@ impl LocalPartitionWriter { Ok(Self { partition_offsets, data_output, + buffers: None, codec_context: ShuffleCodecContext::default(), offsets: vec![0u64; num_output_partitions + 1], batch_size, @@ -169,6 +177,112 @@ impl LocalPartitionWriter { DataOutput::Single { .. } => panic!("single-partition output does not spill"), } } + + pub(crate) fn buffer_bytes_held(&self) -> usize { + match &self.data_output { + DataOutput::Multi { + output_writer, + spill, + spill_reader, + recycled_buffer, + .. + } => { + output_writer.as_ref().map_or(0, BufWriter::capacity) + + spill.buffer_capacity() + + spill_reader.as_ref().map_or(0, |(_, buffer)| buffer.len()) + + recycled_buffer.capacity() + + self.codec_context.retained_bytes() + } + DataOutput::Single { .. } => 0, + } + } + + /// No pool call when nothing changed, so steady-state writes stay off the pool. + fn sync_buffer_reservation(&self) { + if let Some(buffers) = &self.buffers { + buffers.resize(self.buffer_bytes_held()); + } + } + + /// Runs before the first partition is finished. Nothing is spilled after that, so the + /// spill file's write buffer goes back to the pool before the output's buffers are sized: + /// the spill reader first, since it serves every spilled range, then the data file buffer, + /// which blocks of the write buffer size bypass anyway. + fn begin_output( + &mut self, + metrics: &ShufflePartitionerMetrics, + ) -> datafusion::common::Result<()> { + let DataOutput::Multi { spill, .. } = &mut self.data_output else { + return Ok(()); + }; + let mut flush_timer = metrics.write_time.timer(); + spill.release_buffer()?; + flush_timer.stop(); + self.sync_buffer_reservation(); + + let DataOutput::Multi { + output_writer, + spill, + spill_reader, + .. + } = &mut self.data_output + else { + return Ok(()); + }; + let write_buffer_size = self.write_buffer_size; + let reserve = || match self.buffers.as_ref() { + Some(buffers) => super::reserve_buffer(buffers, write_buffer_size), + None => Ok(write_buffer_size), + }; + if spill.has_ranges() { + let path = spill.path()?.ok_or_else(|| { + DataFusionError::Internal( + "shuffle spill ranges recorded without a spill file".to_string(), + ) + })?; + // Open first, so a failed open leaves nothing charged. + let spill_file = File::open(path)?; + *spill_reader = Some((spill_file, vec![0; reserve()?])); + } + let capacity = reserve()?; + let writer = output_writer.take().ok_or_else(missing_output_writer)?; + // An unbuffered writer holds no bytes, so only the file carries over. + let (output_file, _) = writer.into_parts(); + *output_writer = Some(BufWriter::with_capacity(capacity, output_file)); + Ok(()) + } + + #[cfg(test)] + pub(crate) fn buffer_reservation_size(&self) -> usize { + self.buffers.as_ref().map_or(0, MemoryReservation::size) + } + + #[cfg(test)] + pub(crate) fn scratch_capacity(&self) -> usize { + match &self.data_output { + DataOutput::Multi { + recycled_buffer, .. + } => recycled_buffer.capacity(), + DataOutput::Single { scratch, .. } => scratch.capacity(), + } + } + + #[cfg(test)] + pub(crate) fn buffer_capacities(&self) -> (usize, usize, usize) { + match &self.data_output { + DataOutput::Multi { + output_writer, + spill, + spill_reader, + .. + } => ( + output_writer.as_ref().map_or(0, BufWriter::capacity), + spill.buffer_capacity(), + spill_reader.as_ref().map_or(0, |(_, buffer)| buffer.len()), + ), + DataOutput::Single { .. } => panic!("single-partition output has no spill buffers"), + } + } } impl PartitionWriter for LocalPartitionWriter { @@ -181,49 +295,54 @@ impl PartitionWriter for LocalPartitionWriter { where I: Iterator>, { - match &mut self.data_output { - DataOutput::Single { writer, scratch } => { - if pid != 0 { - return Err(DataFusionError::Execution( - "LocalPartitionWriter single-partition output only supports partition 0." - .to_string(), - )); - } + // Synced on every exit, errors included, so it matches what a failed write left. + let result: datafusion::common::Result<()> = (|| { + match &mut self.data_output { + DataOutput::Single { writer, scratch } => { + if pid != 0 { + return Err(DataFusionError::Execution( + "LocalPartitionWriter single-partition output only supports partition 0." + .to_string(), + )); + } - // Stream batches through the long-lived writer so small batches keep - // coalescing across calls. Do not flush here: flushing also finalizes any - // partially coalesced batch, which would defeat cross-call coalescing and - // increase flush frequency. The single-partition writer is flushed once, in - // `finish_all`. - for batch in iter.by_ref() { - let batch = batch?; - writer.write( - &batch, - scratch, - &mut self.codec_context, - &metrics.encode_time, - &metrics.write_time, - )?; + // Stream batches through the long-lived writer so small batches keep + // coalescing across calls. Do not flush here: flushing also finalizes any + // partially coalesced batch, which would defeat cross-call coalescing and + // increase flush frequency. The single-partition writer is flushed once, in + // `finish_all`. + for batch in iter.by_ref() { + let batch = batch?; + writer.write( + &batch, + scratch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; + } } - } - DataOutput::Multi { - spill, - runtime, - recycled_buffer, - .. - } => { - spill.write( - pid, - iter, - &mut self.codec_context, + DataOutput::Multi { + spill, runtime, - metrics, recycled_buffer, - )?; + .. + } => { + spill.write( + pid, + iter, + &mut self.codec_context, + runtime, + metrics, + recycled_buffer, + self.buffers.as_ref(), + )?; + } } - } - - Ok(()) + Ok(()) + })(); + self.sync_buffer_reservation(); + result } fn finish_partition( @@ -240,93 +359,96 @@ impl PartitionWriter for LocalPartitionWriter { "LocalPartitionWriter::finish_partition must be called in order.".to_string(), )); } + let is_first = self.last_finish_pid == -1; self.last_finish_pid = pid as i32; let write_buffer_size = self.write_buffer_size; let batch_size = self.batch_size; - match &mut self.data_output { - DataOutput::Single { writer, scratch } => { - // Single-partition data was already streamed via `write`, starting at - // offset 0 (already recorded in `self.offsets[0]`). Stream any trailing - // batches (normally none) without flushing; the long-lived writer is - // flushed once in `finish_all`. - for batch in iter.by_ref() { - let batch = batch?; - writer.write( - &batch, - scratch, - &mut self.codec_context, - &metrics.encode_time, - &metrics.write_time, - )?; - } + // Synced on every exit, errors included, so it matches what a failed partition left. + let result: datafusion::common::Result<()> = (|| { + if is_first { + self.begin_output(metrics)?; } - DataOutput::Multi { - output_writer, - shuffle_block_writer, - spill, - spill_reader, - recycled_buffer, - .. - } => { - self.offsets[pid] = output_writer.stream_position()?; - - let mut flush_timer = metrics.write_time.timer(); - spill.flush()?; - flush_timer.stop(); - let ranges = spill.ranges(pid)?; - if !ranges.is_empty() { - if spill_reader.is_none() { - let path = spill.path()?.ok_or_else(|| { + match &mut self.data_output { + DataOutput::Single { writer, scratch } => { + // Single-partition data was already streamed via `write`, starting at + // offset 0 (already recorded in `self.offsets[0]`). Stream any trailing + // batches (normally none) without flushing; the long-lived writer is + // flushed once in `finish_all`. + for batch in iter.by_ref() { + let batch = batch?; + writer.write( + &batch, + scratch, + &mut self.codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; + } + } + DataOutput::Multi { + output_writer, + shuffle_block_writer, + spill, + spill_reader, + recycled_buffer, + .. + } => { + let output_writer = output_writer.as_mut().ok_or_else(missing_output_writer)?; + self.offsets[pid] = output_writer.stream_position()?; + + let ranges = spill.ranges(pid)?; + if !ranges.is_empty() { + let (spill_file, buffer) = spill_reader.as_mut().ok_or_else(|| { DataFusionError::Internal( - "shuffle spill ranges recorded without a spill file".to_string(), + "shuffle spill ranges recorded without a spill reader".to_string(), ) })?; - *spill_reader = Some((File::open(path)?, vec![0; write_buffer_size])); - } - let (spill_file, buffer) = spill_reader.as_mut().unwrap(); - let mut write_timer = metrics.write_time.timer(); - for range in ranges { - copy_spill_range(spill_file, buffer, range, output_writer)?; + let mut write_timer = metrics.write_time.timer(); + for range in ranges { + copy_spill_range(spill_file, buffer, range, output_writer)?; + } + write_timer.stop(); } - write_timer.stop(); - } - // Write in memory batches to output data file. Each partition uses its - // own writer so coalescing does not cross partition boundaries, but the - // scratch buffer is shared so its capacity carries over to the next one. - let mut buf_batch_writer = BufBatchWriter::new( - shuffle_block_writer, - output_writer, - write_buffer_size, - batch_size, - ); - let codec_context = &mut self.codec_context; - let result: datafusion::common::Result<()> = (|| { - for batch in iter.by_ref() { - let batch = batch?; - buf_batch_writer.write( - &batch, + // Write in memory batches to output data file. Each partition uses its + // own writer so coalescing does not cross partition boundaries, but the + // scratch buffer is shared so its capacity carries over to the next one. + let mut buf_batch_writer = BufBatchWriter::new( + shuffle_block_writer, + output_writer, + write_buffer_size, + batch_size, + ); + let codec_context = &mut self.codec_context; + let result: datafusion::common::Result<()> = (|| { + for batch in iter.by_ref() { + let batch = batch?; + buf_batch_writer.write( + &batch, + recycled_buffer, + codec_context, + &metrics.encode_time, + &metrics.write_time, + )?; + } + buf_batch_writer.flush( recycled_buffer, codec_context, &metrics.encode_time, &metrics.write_time, - )?; - } - buf_batch_writer.flush( - recycled_buffer, - codec_context, - &metrics.encode_time, - &metrics.write_time, - ) - })(); - // An errored partition must hand back a drained buffer, or its bytes - // leak into the next partition's block. - result.inspect_err(|_| recycled_buffer.clear())?; + ) + })(); + // An errored partition must hand back a drained buffer, or its bytes + // leak into the next partition's block. + result.inspect_err(|_| recycled_buffer.clear())?; + } } - } - Ok(()) + Ok(()) + })(); + self.sync_buffer_reservation(); + result } fn finish_all( @@ -346,6 +468,7 @@ impl PartitionWriter for LocalPartitionWriter { writer.writer_stream_position()? } DataOutput::Multi { output_writer, .. } => { + let output_writer = output_writer.as_mut().ok_or_else(missing_output_writer)?; let mut write_timer = metrics.write_time.timer(); output_writer.flush()?; let pos = output_writer.stream_position()?; @@ -372,15 +495,51 @@ impl PartitionWriter for LocalPartitionWriter { // The shuffle output is complete; nothing else encodes through this context. self.codec_context.release_zstd(); + self.sync_buffer_reservation(); Ok(()) } - fn write_burst_complete(&mut self) { - // A spill burst just ended and the next encode may be a long time coming; the zstd - // workspace is native memory no reservation tracks, so don't sit on it. + fn write_burst_complete(&mut self, freed: usize) -> datafusion::common::Result<()> { + // A spill burst just ended and the next encode may be a long time coming, so don't + // sit on the zstd workspace, nor on the drained scratch, which one wide row can leave + // larger than the pool. self.codec_context.release_zstd(); + if let DataOutput::Multi { + recycled_buffer, .. + } = &mut self.data_output + { + *recycled_buffer = Vec::new(); + } + self.sync_buffer_reservation(); + // A pressure spill opened its file while the batches still held the pool; now they don't. + // Half of what the spill freed stays with the input, so the next batch does not spill. + let capacity = self.write_buffer_size.min(freed / 2); + if let (Some(buffers), DataOutput::Multi { spill, .. }) = + (&self.buffers, &mut self.data_output) + { + spill.restore_buffer(buffers, capacity)?; + } + Ok(()) } + + fn attach_buffer_reservation( + &mut self, + reservation: MemoryReservation, + ) -> datafusion::common::Result<()> { + let DataOutput::Multi { .. } = &self.data_output else { + return Ok(()); + }; + // Attached once, before any write. Each buffer is charged when it is allocated. + debug_assert!(self.buffers.is_none()); + self.buffers = Some(reservation); + self.sync_buffer_reservation(); + Ok(()) + } +} + +fn missing_output_writer() -> DataFusionError { + DataFusionError::Internal("shuffle data file writer missing".to_string()) } /// Appends `range` of the spill file to `output`, reading it through `buffer` when it fits. @@ -431,14 +590,19 @@ mod tests { use arrow::compute::concat_batches; use arrow::datatypes::{DataType, Field, Schema}; use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; fn test_batch() -> RecordBatch { + test_batch_of(100) + } + + fn test_batch_of(rows: i64) -> RecordBatch { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); RecordBatch::try_new( Arc::clone(&schema), - vec![Arc::new(Int64Array::from_iter_values(0..100))], + vec![Arc::new(Int64Array::from_iter_values(0..rows))], ) .unwrap() } @@ -458,9 +622,25 @@ mod tests { dir: &tempfile::TempDir, runtime: Arc, ) -> LocalPartitionWriter { - let block_writer = - ShuffleBlockWriter::try_new(batch.schema_ref().as_ref(), CompressionCodec::None) - .unwrap(); + partition_writer_with_codec( + batch, + num_partitions, + write_buffer_size, + CompressionCodec::None, + dir, + runtime, + ) + } + + fn partition_writer_with_codec( + batch: &RecordBatch, + num_partitions: usize, + write_buffer_size: usize, + codec: CompressionCodec, + dir: &tempfile::TempDir, + runtime: Arc, + ) -> LocalPartitionWriter { + let block_writer = ShuffleBlockWriter::try_new(batch.schema_ref().as_ref(), codec).unwrap(); LocalPartitionWriter::try_new( dir.path().join("data.out").to_str().unwrap().to_string(), Arc::new(PartitionOffsets::default()), @@ -681,4 +861,243 @@ mod tests { ); } } + + fn greedy_pool(limit: usize) -> Arc { + Arc::new(GreedyMemoryPool::new(limit)) + } + + fn assert_buffers_charged(writer: &LocalPartitionWriter, pool: &Arc) { + assert_eq!( + writer.buffer_reservation_size(), + writer.buffer_bytes_held(), + "reservation must equal the bytes held" + ); + assert_eq!(pool.reserved(), writer.buffer_bytes_held()); + } + + /// Spills partition 0's 4096-row block and partition 1's 100-row block, then writes the + /// output under a pool of `limit` bytes. Returns the buffer capacities, the output, and the + /// encode scratch left by the spill. + fn write_spilled_output(limit: usize) -> ((usize, usize, usize), Vec, usize) { + let write_buffer_size = 1 << 20; + let fallback = crate::writers::local::FALLBACK_BUFFER_SIZE; + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + // Partition 0's block exceeds the fallback buffer (`io::copy`), partition 1's fits (pread). + let batch_for = |pid: usize| { + if pid == 0 { + test_batch_of(4096) + } else { + test_batch() + } + }; + let dir = tempfile::tempdir().unwrap(); + let pool = greedy_pool(limit); + let mut writer = partition_writer_with( + &test_batch(), + 2, + write_buffer_size, + &dir, + Arc::new(RuntimeEnv::default()), + ); + let reservation = MemoryConsumer::new("test").register(&pool); + writer + .attach_buffer_reservation(reservation.new_empty()) + .unwrap(); + for pid in 0..2 { + writer + .write(pid, &mut vec![Ok(batch_for(pid))].into_iter(), &metrics) + .unwrap(); + } + let scratch = writer.scratch_capacity(); + writer.write_burst_complete(0).unwrap(); + let range_len = |pid: usize| { + let ranges = writer.get_spill().ranges(pid).unwrap(); + assert_eq!(ranges.len(), 1, "one spilled block for partition {pid}"); + (ranges[0].end - ranges[0].start) as usize + }; + assert!( + range_len(0) > fallback, + "partition 0 must exceed the fallback buffer" + ); + assert!( + range_len(1) <= fallback, + "partition 1 must fit the fallback buffer" + ); + for pid in 0..2 { + writer + .finish_partition(pid, &mut vec![Ok(batch_for(pid))].into_iter(), &metrics) + .unwrap(); + } + writer.finish_all(&metrics).unwrap(); + assert_buffers_charged(&writer, &pool); + ( + writer.buffer_capacities(), + std::fs::read(dir.path().join("data.out")).unwrap(), + scratch, + ) + } + + /// The output buffers are sized once nothing more is spilled, after the spill file's write + /// buffer is released, and the reader is sized before the data file buffer. A refused + /// reader falls back to a small buffer and copies larger ranges with `io::copy`. + #[test] + fn output_buffers_are_sized_after_the_spill_buffer_is_released() { + let write_buffer_size = 1 << 20; + let fallback = crate::writers::local::FALLBACK_BUFFER_SIZE; + let (capacities, baseline, scratch) = write_spilled_output(1 << 30); + assert_eq!(capacities, (write_buffer_size, 0, write_buffer_size)); + assert!(!baseline.is_empty()); + for (limit, expected) in [ + // room for two full buffers only once the spill buffer is gone + ( + 2 * write_buffer_size + scratch + 64 * 1024, + (write_buffer_size, 0, write_buffer_size), + ), + // room for one full buffer, which goes to the reader + ( + write_buffer_size + scratch + 64 * 1024, + (fallback, 0, write_buffer_size), + ), + (64 * 1024, (fallback, 0, fallback)), + ] { + let (capacities, output, _) = write_spilled_output(limit); + assert_eq!(capacities, expected, "pool limit {limit}"); + assert_eq!(output, baseline, "pool limit {limit} changed the output"); + } + } + + /// A partition that fails after the spill buffer was released still leaves the + /// reservation equal to what the writer holds. + #[test] + fn failed_first_partition_keeps_the_reservation_in_sync() { + let dir = tempfile::tempdir().unwrap(); + let pool = greedy_pool(1 << 30); + let mut writer = partition_writer_with( + &test_batch(), + 2, + 1 << 20, + &dir, + Arc::new(RuntimeEnv::default()), + ); + let reservation = MemoryConsumer::new("test").register(&pool); + writer + .attach_buffer_reservation(reservation.new_empty()) + .unwrap(); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + writer + .write(0, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); + assert_eq!(writer.buffer_capacities().1, 1 << 20); + let scratch = writer.scratch_capacity(); + + // The larger batch grows the scratch before the iterator fails. + let mut iter = vec![ + Ok(test_batch_of(4096)), + Err(DataFusionError::Execution("injected failure".to_string())), + ] + .into_iter(); + assert!(writer.finish_partition(0, &mut iter, &metrics).is_err()); + assert!(writer.scratch_capacity() > scratch); + assert_eq!(writer.buffer_capacities().1, 0, "spill buffer released"); + assert_buffers_charged(&writer, &pool); + } + + #[test] + fn encode_scratch_growth_is_charged() { + let dir = tempfile::tempdir().unwrap(); + let pool = greedy_pool(1 << 30); + let mut writer = partition_writer_with( + &test_batch(), + 2, + 1 << 20, + &dir, + Arc::new(RuntimeEnv::default()), + ); + let reservation = MemoryConsumer::new("test").register(&pool); + writer + .attach_buffer_reservation(reservation.new_empty()) + .unwrap(); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + + writer + .write(0, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); + let (charged, scratch) = (writer.buffer_reservation_size(), writer.scratch_capacity()); + + // A larger block grows only the scratch; the file buffers and codec are unchanged. + writer + .write(1, &mut vec![Ok(test_batch_of(4096))].into_iter(), &metrics) + .unwrap(); + let scratch_delta = writer.scratch_capacity() - scratch; + assert!(scratch_delta > 0, "the larger block must grow the scratch"); + assert_eq!( + writer.buffer_reservation_size() - charged, + scratch_delta, + "the reservation must move by the scratch's growth" + ); + assert_buffers_charged(&writer, &pool); + } + + #[test] + fn zstd_context_is_charged_while_retained() { + let dir = tempfile::tempdir().unwrap(); + let pool = greedy_pool(1 << 30); + let mut writer = partition_writer_with_codec( + &test_batch(), + 2, + 1 << 20, + CompressionCodec::Zstd(1), + &dir, + Arc::new(RuntimeEnv::default()), + ); + let reservation = MemoryConsumer::new("test").register(&pool); + writer + .attach_buffer_reservation(reservation.new_empty()) + .unwrap(); + let metrics = ShufflePartitionerMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + + writer + .write(0, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); + let retained = writer.codec_context.retained_bytes(); + assert!( + retained > 1 << 20, + "level 1 context is over 1 MiB, got {retained}" + ); + assert_buffers_charged(&writer, &pool); + let after_write = writer.buffer_reservation_size(); + + let scratch = writer.scratch_capacity(); + writer.write_burst_complete(0).unwrap(); + assert!(!writer.holds_zstd_cctx()); + assert_eq!( + writer.buffer_reservation_size(), + after_write - retained - scratch + ); + assert_buffers_charged(&writer, &pool); + + for pid in 0..2 { + writer + .finish_partition(pid, &mut vec![Ok(test_batch())].into_iter(), &metrics) + .unwrap(); + } + let retained = writer.codec_context.retained_bytes(); + assert!( + retained > 1 << 20, + "finish_partition re-created the context" + ); + assert_buffers_charged(&writer, &pool); + let after_finish_partition = writer.buffer_reservation_size(); + + writer.finish_all(&metrics).unwrap(); + assert!(!writer.holds_zstd_cctx()); + assert_eq!( + writer.buffer_reservation_size(), + after_finish_partition - retained + ); + assert_buffers_charged(&writer, &pool); + + drop(writer); + assert_eq!(pool.reserved(), 0); + } } diff --git a/native/shuffle/src/writers/local/mod.rs b/native/shuffle/src/writers/local/mod.rs index e76bfba21a2..de37e569417 100644 --- a/native/shuffle/src/writers/local/mod.rs +++ b/native/shuffle/src/writers/local/mod.rs @@ -17,3 +17,134 @@ pub(crate) mod local_partition_writer; mod spill; + +use datafusion::common::DataFusionError; +use datafusion::execution::memory_pool::MemoryReservation; + +const FALLBACK_BUFFER_SIZE: usize = 8 * 1024; + +/// Whether `error` is the pool refusing memory. Other errors, such as an exception Spark threw +/// while spilling another consumer, must fail the task. +fn is_refusal(error: &DataFusionError) -> bool { + matches!(error.find_root(), DataFusionError::ResourcesExhausted(_)) +} + +/// Returns the capacity to allocate. A refused size falls back to a small buffer charged with +/// `grow`, so the task carries on and the reservation still matches what is held. +fn reserve_buffer( + reservation: &MemoryReservation, + size: usize, +) -> datafusion::common::Result { + match reservation.try_grow(size) { + Ok(()) => Ok(size), + Err(error) if is_refusal(&error) => { + let fallback = size.min(FALLBACK_BUFFER_SIZE); + reservation.grow(fallback); + Ok(fallback) + } + Err(error) => Err(error), + } +} + +#[cfg(test)] +pub(crate) mod erroring_pool { + use datafusion::common::DataFusionError; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, MemoryReservation}; + use std::fmt; + + /// A greedy pool whose `try_grow` fails with `error()` for requests of `fail_at` bytes or more. + #[derive(Debug)] + pub(crate) struct ErroringPool { + inner: GreedyMemoryPool, + fail_at: usize, + error: fn() -> DataFusionError, + } + + impl ErroringPool { + pub(crate) fn new(limit: usize, fail_at: usize, error: fn() -> DataFusionError) -> Self { + Self { + inner: GreedyMemoryPool::new(limit), + fail_at, + error, + } + } + } + + /// What a Spark-backed pool returns when Spark throws. + pub(crate) fn external() -> DataFusionError { + DataFusionError::External(Box::new(std::io::Error::other("spill of a peer failed"))) + } + + /// A refusal under a context, so the check has to look at the root. + pub(crate) fn refusal() -> DataFusionError { + DataFusionError::ResourcesExhausted("refused".to_string()).context("wrapped") + } + + impl fmt::Display for ErroringPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "erroring") + } + } + + impl MemoryPool for ErroringPool { + fn name(&self) -> &str { + "erroring" + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional) + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink) + } + + fn try_grow( + &self, + reservation: &MemoryReservation, + additional: usize, + ) -> datafusion::common::Result<()> { + if additional >= self.fail_at { + return Err((self.error)()); + } + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + } +} + +#[cfg(test)] +mod tests { + use super::erroring_pool::{external, refusal, ErroringPool}; + use super::*; + use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool}; + use std::sync::Arc; + + #[test] + fn reserve_buffer_falls_back_only_on_refusal() { + let size = 1 << 20; + let pool: Arc = Arc::new(ErroringPool::new(1 << 30, size, refusal)); + let reservation = MemoryConsumer::new("test").register(&pool); + assert_eq!( + reserve_buffer(&reservation, size).unwrap(), + FALLBACK_BUFFER_SIZE + ); + assert_eq!(reservation.size(), FALLBACK_BUFFER_SIZE); + } + + #[test] + fn reserve_buffer_propagates_other_pool_errors() { + let size = 1 << 20; + let pool: Arc = Arc::new(ErroringPool::new(1 << 30, size, external)); + let reservation = MemoryConsumer::new("test").register(&pool); + let error = reserve_buffer(&reservation, size).expect_err("an exception is not a refusal"); + assert!( + matches!(error, DataFusionError::External(_)), + "unexpected error: {error}" + ); + assert_eq!((reservation.size(), pool.reserved()), (0, 0)); + } +} diff --git a/native/shuffle/src/writers/local/spill.rs b/native/shuffle/src/writers/local/spill.rs index ae4b8adfa44..8f355f426d7 100644 --- a/native/shuffle/src/writers/local/spill.rs +++ b/native/shuffle/src/writers/local/spill.rs @@ -21,6 +21,7 @@ use crate::writers::BufBatchWriter; use crate::ShuffleBlockWriter; use arrow::record_batch::RecordBatch; use datafusion::common::DataFusionError; +use datafusion::execution::memory_pool::MemoryReservation; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::execution::SpillFile as DfSpillFile; use datafusion::execution::SpillWriter as DfSpillWriter; @@ -88,6 +89,7 @@ impl PartitionedSpill { /// Appends partition `pid`'s batches to the spill file. `recycled_buffer` is left drained, /// including on error. + #[allow(clippy::too_many_arguments)] pub(crate) fn write>>( &mut self, pid: usize, @@ -96,12 +98,13 @@ impl PartitionedSpill { runtime: &RuntimeEnv, metrics: &ShufflePartitionerMetrics, recycled_buffer: &mut Vec, + buffers: Option<&MemoryReservation>, ) -> datafusion::common::Result<()> { self.check_usable()?; let Some(batch) = iter.next() else { return Ok(()); }; - self.ensure_spill_file_created(runtime)?; + self.ensure_spill_file_created(runtime, buffers)?; let result = (|| { let mut buf_batch_writer = BufBatchWriter::new( @@ -214,17 +217,87 @@ impl PartitionedSpill { fn ensure_spill_file_created( &mut self, runtime: &RuntimeEnv, + buffers: Option<&MemoryReservation>, ) -> datafusion::common::Result<()> { if self.spill_file.is_none() { let temp_file = runtime .disk_manager .create_tmp_file("shuffle writer spill")?; - let writer = BufWriter::with_capacity(self.write_buffer_size, temp_file.open_writer()?); + // Open before charging, so a failed open leaves nothing on the reservation. + let file = temp_file.open_writer()?; + let capacity = match buffers { + Some(reservation) => super::reserve_buffer(reservation, self.write_buffer_size)?, + None => self.write_buffer_size, + }; + let writer = BufWriter::with_capacity(capacity, file); self.spill_file = Some(ActiveSpillFile { temp_file, writer }); } Ok(()) } + /// Grows the write buffer to `capacity` if `buffers` grants it, and never shrinks it. Buffered + /// bytes move to the new buffer; nothing is written to the file. + /// Errors other than a refusal are returned. + pub(crate) fn restore_buffer( + &mut self, + buffers: &MemoryReservation, + capacity: usize, + ) -> datafusion::common::Result<()> { + let current = self.buffer_capacity(); + if self.failed || self.spill_file.is_none() || capacity <= current { + return Ok(()); + } + match buffers.try_grow(capacity - current) { + Ok(()) => { + self.resize_buffer(capacity); + Ok(()) + } + Err(error) if super::is_refusal(&error) => Ok(()), + Err(error) => Err(error), + } + } + + /// Writes the buffered bytes to the file and drops the write buffer, once nothing more will + /// be spilled. + pub(crate) fn release_buffer(&mut self) -> datafusion::common::Result<()> { + self.flush()?; + self.resize_buffer(0); + Ok(()) + } + + /// Rebuilds the write buffer at `capacity`, moving the buffered bytes into it. Callers make + /// sure they fit, so nothing is written to the file. + fn resize_buffer(&mut self, capacity: usize) { + let Some(ActiveSpillFile { temp_file, writer }) = self.spill_file.take() else { + return; + }; + let (file, buffered) = writer.into_parts(); + let buffered = match buffered { + Ok(buffered) => buffered, + Err(panicked) => { + // a write panicked partway, so the file may already hold some of these bytes + self.failed = true; + panicked.into_inner() + } + }; + let mut writer = BufWriter::with_capacity(capacity, file); + if writer.write_all(&buffered).is_err() { + self.failed = true; + } + self.spill_file = Some(ActiveSpillFile { temp_file, writer }); + } + + /// Whether any partition has spilled bytes. + pub(crate) fn has_ranges(&self) -> bool { + self.len > 0 + } + + pub(crate) fn buffer_capacity(&self) -> usize { + self.spill_file + .as_ref() + .map_or(0, |spill_file| spill_file.writer.capacity()) + } + #[cfg(test)] pub(crate) fn has_spill_file(&self) -> bool { self.spill_file.is_some() @@ -318,6 +391,7 @@ mod tests { use crate::CompressionCodec; use arrow::array::Int64Array; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool}; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use std::sync::Arc; @@ -355,7 +429,8 @@ mod tests { &mut ShuffleCodecContext::default(), &RuntimeEnv::default(), &metrics(), - recycled + recycled, + None ) .is_err()); } @@ -386,6 +461,7 @@ mod tests { &RuntimeEnv::default(), &metrics(), &mut recycled, + None, ) .expect_err("write after a failed write"); assert!( @@ -410,6 +486,7 @@ mod tests { &runtime, &metrics(), &mut recycled, + None, ) .unwrap(); } @@ -438,6 +515,7 @@ mod tests { &runtime, &metrics(), &mut recycled, + None, ) .unwrap(); } @@ -449,6 +527,134 @@ mod tests { assert_eq!(std::fs::metadata(&path).unwrap().len(), spilled); } + /// `partitioned_spill`'s configured write buffer size. + const WRITE_BUFFER: usize = 1 << 20; + + fn full_pool() -> (Arc, MemoryReservation, MemoryReservation) { + let pool: Arc = Arc::new(GreedyMemoryPool::new(WRITE_BUFFER)); + let batches = MemoryConsumer::new("batches").register(&pool); + batches.grow(WRITE_BUFFER); + let buffers = MemoryConsumer::new("buffers").register(&pool); + (pool, batches, buffers) + } + + fn write_with_buffers>>( + spill: &mut PartitionedSpill, + iter: &mut I, + buffers: &MemoryReservation, + ) -> datafusion::common::Result<()> { + spill.write( + 0, + iter, + &mut ShuffleCodecContext::default(), + &RuntimeEnv::default(), + &metrics(), + &mut Vec::new(), + Some(buffers), + ) + } + + #[test] + fn restore_buffer_moves_buffered_bytes_without_writing() { + let fallback = super::super::FALLBACK_BUFFER_SIZE; + let (pool, batches, buffers) = full_pool(); + let mut spill = partitioned_spill(&test_batch(), 1); + write_with_buffers( + &mut spill, + &mut vec![Ok(test_batch())].into_iter(), + &buffers, + ) + .unwrap(); + assert_eq!( + (spill.buffer_capacity(), buffers.size()), + (fallback, fallback) + ); + let path = spill.path().unwrap().unwrap().to_path_buf(); + let spilled = spill.ranges(0).unwrap()[0].end; + assert!( + spilled > 0 && spilled < fallback as u64, + "the bytes must sit in the buffer" + ); + + spill.restore_buffer(&buffers, WRITE_BUFFER).unwrap(); + assert_eq!( + (spill.buffer_capacity(), buffers.size()), + (fallback, fallback), + "refused while the batches hold the pool" + ); + + batches.free(); + spill.restore_buffer(&buffers, WRITE_BUFFER).unwrap(); + assert_eq!( + (spill.buffer_capacity(), buffers.size(), pool.reserved()), + (WRITE_BUFFER, WRITE_BUFFER, WRITE_BUFFER) + ); + assert_eq!(std::fs::metadata(&path).unwrap().len(), 0); + + spill.flush().unwrap(); + let bytes = std::fs::read(&path).unwrap(); + assert_eq!(bytes.len() as u64, spilled); + let mut values = Vec::new(); + let mut pos = 0; + while pos < bytes.len() { + let len = u64::from_le_bytes(bytes[pos..pos + 8].try_into().unwrap()) as usize; + let batch = crate::read_ipc_compressed(&bytes[pos + 16..pos + 8 + len]).unwrap(); + let column = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + values.extend_from_slice(column.values()); + pos += 8 + len; + } + assert_eq!(values, (0..100).collect::>()); + } + + #[test] + fn restore_buffer_returns_errors_other_than_refusal() { + use super::super::erroring_pool::{external, refusal, ErroringPool}; + let fallback = super::super::FALLBACK_BUFFER_SIZE; + let (_pool, _batches, buffers) = full_pool(); + let mut spill = partitioned_spill(&test_batch(), 1); + write_with_buffers( + &mut spill, + &mut vec![Ok(test_batch())].into_iter(), + &buffers, + ) + .unwrap(); + + for (error, is_refusal) in [ + (refusal as fn() -> DataFusionError, true), + (external, false), + ] { + let pool: Arc = Arc::new(ErroringPool::new(1 << 30, 0, error)); + let reservation = MemoryConsumer::new("buffers").register(&pool); + let result = spill.restore_buffer(&reservation, WRITE_BUFFER); + assert_eq!(result.is_ok(), is_refusal, "{result:?}"); + assert_eq!((spill.buffer_capacity(), reservation.size()), (fallback, 0)); + } + } + + #[test] + fn restore_buffer_skips_a_failed_spill_file() { + let fallback = super::super::FALLBACK_BUFFER_SIZE; + let (_pool, batches, buffers) = full_pool(); + let mut spill = partitioned_spill(&test_batch(), 1); + let mut iter = vec![ + Ok(test_batch()), + Err(DataFusionError::Execution("injected failure".to_string())), + ] + .into_iter(); + assert!(write_with_buffers(&mut spill, &mut iter, &buffers).is_err()); + + batches.free(); + spill.restore_buffer(&buffers, WRITE_BUFFER).unwrap(); + assert_eq!( + (spill.buffer_capacity(), buffers.size()), + (fallback, fallback) + ); + } + #[test] fn empty_write_records_no_range() { let mut spill = partitioned_spill(&test_batch(), 2); @@ -460,6 +666,7 @@ mod tests { &RuntimeEnv::default(), &metrics(), &mut Vec::new(), + None, ) .unwrap(); assert!(!spill.has_spill_file()); @@ -484,6 +691,7 @@ mod tests { &pathless_backend::runtime(), &metrics(), &mut Vec::new(), + None, ) .unwrap(); assert!(spill.has_spill_file()); diff --git a/native/shuffle/src/writers/mod.rs b/native/shuffle/src/writers/mod.rs index 4586e46c25b..75cd5db1394 100644 --- a/native/shuffle/src/writers/mod.rs +++ b/native/shuffle/src/writers/mod.rs @@ -24,6 +24,8 @@ mod shuffle_block_writer; pub(crate) use buf_batch_writer::BufBatchWriter; pub(crate) use checksum::Checksum; +#[cfg(test)] +pub(crate) use local::erroring_pool; pub(crate) use local::local_partition_writer::LocalPartitionWriter; pub(crate) use partition_writer::PartitionWriter; pub(crate) use rss::rss_partition_writer::RssPartitionWriter; diff --git a/native/shuffle/src/writers/partition_writer.rs b/native/shuffle/src/writers/partition_writer.rs index 0fcec17c00d..683d37e87cb 100644 --- a/native/shuffle/src/writers/partition_writer.rs +++ b/native/shuffle/src/writers/partition_writer.rs @@ -17,6 +17,7 @@ use crate::metrics::ShufflePartitionerMetrics; use arrow::record_batch::RecordBatch; +use datafusion::execution::memory_pool::MemoryReservation; /// Storage backend abstraction for shuffle partition output. /// @@ -72,5 +73,17 @@ pub(crate) trait PartitionWriter: Send { /// Marks the end of one burst of [`write`](PartitionWriter::write) calls (a spill /// event), letting the writer drop transient encode state. Staging more batches /// afterwards is still allowed. - fn write_burst_complete(&mut self) {} + /// Called after the spilled batches are freed, so the writer can ask the pool again; + /// `freed` is the number of bytes the spill released. + fn write_burst_complete(&mut self, _freed: usize) -> datafusion::common::Result<()> { + Ok(()) + } + + /// Hands the writer an empty reservation for its own buffers. Called once, before any `write`. + fn attach_buffer_reservation( + &mut self, + _reservation: MemoryReservation, + ) -> datafusion::common::Result<()> { + Ok(()) + } } diff --git a/native/shuffle/src/writers/shuffle_block_writer.rs b/native/shuffle/src/writers/shuffle_block_writer.rs index dc1bbf40dc9..0a888b70c96 100644 --- a/native/shuffle/src/writers/shuffle_block_writer.rs +++ b/native/shuffle/src/writers/shuffle_block_writer.rs @@ -361,8 +361,8 @@ impl ShuffleBlockWriter { // longer tracks. codec_context.release_zstd(); } else { - // Local shuffle reuses the context across blocks, but nothing reserves its - // memory: a high-level workspace (hundreds of MiB) must not outlive the block. + // Local shuffle reuses the context across blocks and charges only what it keeps: + // a workspace over the retention cap (high levels) must not outlive the block. codec_context.release_zstd_if_oversized(); } encode_result?;