From 62724f162e52fe55560959d0027d5e0fe1fdc779 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 1 Sep 2026 23:17:57 +0000 Subject: [PATCH 1/2] fix: measure native shuffle spills from materialized batches --- native/shuffle/src/metrics.rs | 4 +- .../src/partitioners/multi_partition.rs | 229 +++++++++++------- native/shuffle/src/shuffle_writer.rs | 86 +++++-- 3 files changed, 217 insertions(+), 102 deletions(-) diff --git a/native/shuffle/src/metrics.rs b/native/shuffle/src/metrics.rs index bcb4439b784..5648f8649f1 100644 --- a/native/shuffle/src/metrics.rs +++ b/native/shuffle/src/metrics.rs @@ -46,7 +46,9 @@ pub(crate) struct ShufflePartitionerMetrics { /// total spilled bytes during the execution of the operator pub(crate) spilled_bytes: Count, - /// Total in-memory bytes released by spills before compression. + /// Total buffer size of materialized spill batches before compression, plus the + /// partition-index allocations released by spills. Measured from spill output rather + /// than input batch boundaries; not a count of globally unique input allocations. pub(crate) memory_spilled_bytes: Count, /// The original size of spilled data. Different to `spilled_bytes` because of compression. diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index 37ce85d1cd2..ab2e4dee1cd 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -21,6 +21,7 @@ use crate::partitioners::ShufflePartitioner; use crate::writers::PartitionWriter; use crate::{comet_partitioning, CometPartitioning}; use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch}; +use datafusion::common::utils::memory::get_record_batch_memory_size; use datafusion::common::utils::proxy::VecAllocExt; use datafusion::common::{DataFusionError, HashSet}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; @@ -122,21 +123,12 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { /// allocation once rather than once per slice that references it. Cleared whenever the /// buffered batches drain (spill / shuffle_write). See `count_new_buffers`. pinned_buffers: HashSet, - /// Backing buffers already reported as spilled while slicing the current outer input batch. - /// The outer batch keeps these allocations alive across spills, so clear this set only when - /// that input batch finishes rather than whenever the repartitioner's buffers drain. - spill_accounted_input_buffers: HashSet, - /// Bytes in the currently buffered batches that were already counted by a previous spill of - /// the same outer input batch. Partition-index allocations are never included here. - repeated_spill_buffer_bytes: usize, } /// Sum of the capacities of the backing buffers reachable from `batch` whose start address is /// not already in `seen` (recursing through child data: dictionary values, list children, and so /// on). `seen` is kept across every buffered batch, so this returns the bytes a batch newly -/// pins, which is the memory the shuffle writer holds resident by buffering it. The second return -/// value contains the subset of those bytes whose buffers were already reported spilled while -/// processing the current outer input batch. +/// pins, which is the memory the shuffle writer holds resident by buffering it. /// /// Cheaper measures do not match resident memory for the batches this writer sees. A partial /// `HashAggregate` emits one group-values buffer sliced into batch_size chunks, and every @@ -153,25 +145,12 @@ pub(crate) struct MultiPartitionShuffleRepartitioner { /// /// Counting each distinct allocation once, keyed by start address, is the measure that tracks /// resident memory regardless of how arrays share or slice their buffers. -fn count_new_buffers( - batch: &RecordBatch, - seen: &mut HashSet, - previously_spilled: Option<&HashSet>, -) -> (usize, usize) { - fn visit( - data: &ArrayData, - seen: &mut HashSet, - previously_spilled: Option<&HashSet>, - total: &mut usize, - repeated: &mut usize, - ) { +fn count_new_buffers(batch: &RecordBatch, seen: &mut HashSet) -> usize { + fn visit(data: &ArrayData, seen: &mut HashSet, total: &mut usize) { for buffer in data.buffers() { let address = buffer.data_ptr().as_ptr() as usize; if seen.insert(address) { *total += buffer.capacity(); - if previously_spilled.is_some_and(|buffers| buffers.contains(&address)) { - *repeated += buffer.capacity(); - } } } if let Some(nulls) = data.nulls() { @@ -179,27 +158,17 @@ fn count_new_buffers( let address = inner.data_ptr().as_ptr() as usize; if seen.insert(address) { *total += inner.capacity(); - if previously_spilled.is_some_and(|buffers| buffers.contains(&address)) { - *repeated += inner.capacity(); - } } } for child in data.child_data() { - visit(child, seen, previously_spilled, total, repeated); + visit(child, seen, total); } } let mut total = 0; - let mut repeated = 0; for column in batch.columns() { - visit( - &column.to_data(), - seen, - previously_spilled, - &mut total, - &mut repeated, - ); + visit(&column.to_data(), seen, &mut total); } - (total, repeated) + total } impl MultiPartitionShuffleRepartitioner { @@ -253,8 +222,6 @@ impl MultiPartitionShuffleRepartitioner { max_buffer_bytes, tracing_enabled, pinned_buffers: HashSet::new(), - spill_accounted_input_buffers: HashSet::new(), - repeated_spill_buffer_bytes: 0, }) } @@ -462,12 +429,7 @@ impl MultiPartitionShuffleRepartitioner { ) -> datafusion::common::Result<()> { // Charge both the reservation and the data_size metric for the buffers this batch newly // pins; `count_new_buffers` dedups buffers shared across already-buffered batches. - let (new_buffer_bytes, repeated_buffer_bytes) = count_new_buffers( - &input, - &mut self.pinned_buffers, - Some(&self.spill_accounted_input_buffers), - ); - self.repeated_spill_buffer_bytes += repeated_buffer_bytes; + let new_buffer_bytes = count_new_buffers(&input, &mut self.pinned_buffers); self.metrics.data_size.add(new_buffer_bytes); let mut mem_growth: usize = new_buffer_bytes; let buffered_partition_idx = self.buffered_batches.len() as u32; @@ -508,15 +470,7 @@ impl MultiPartitionShuffleRepartitioner { .max_buffer_bytes .is_some_and(|limit| self.reservation.size() >= limit) { - let unreserved_bytes = if reservation_failed { mem_growth } else { 0 }; - count_new_buffers( - self.buffered_batches - .last() - .expect("the current input batch was buffered before spilling"), - &mut self.spill_accounted_input_buffers, - None, - ); - self.spill(unreserved_bytes)?; + self.spill()?; } Ok(()) @@ -552,7 +506,7 @@ impl MultiPartitionShuffleRepartitioner { PartitionedBatchesProducer::new(buffered_batches, indices, self.batch_size) } - pub(crate) fn spill(&mut self, unreserved_bytes: usize) -> datafusion::common::Result<()> { + pub(crate) fn spill(&mut self) -> datafusion::common::Result<()> { log::info!( "ShuffleRepartitioner spilling {} bytes to its partition writer ({} previous spills)", self.used(), @@ -566,26 +520,36 @@ impl MultiPartitionShuffleRepartitioner { with_trace("shuffle_spill", self.tracing_enabled, || { let num_output_partitions = self.partition_indices.len(); + // The reservation measures pinned input allocations, not the size of the data + // being spilled: many input batches may be slices of the same large allocation. + // Measure the already-materialized output batches instead, before compression, + // and include the partition-index allocations released by this spill. This is + // cumulative across spills without retaining input buffers or their addresses. + let mut memory_spilled_bytes = self + .partition_indices + .iter() + .map(|indices| indices.allocated_size()) + .sum::(); let write_result = { let mut partitioned_batches = self.partitioned_batches(); (0..num_output_partitions).try_for_each(|partition_id| { - self.partition_writer.write( - partition_id, - &mut partitioned_batches - .produce(partition_id, &self.metrics.interleave_time), - &self.metrics, - ) + let mut batches = partitioned_batches + .produce(partition_id, &self.metrics.interleave_time) + .inspect(|result| { + if let Ok(batch) = result { + memory_spilled_bytes += get_record_batch_memory_size(batch); + } + }); + self.partition_writer + .write(partition_id, &mut batches, &self.metrics) }) }; - let memory_spilled_bytes = self - .reservation - .free() - .saturating_add(unreserved_bytes) - .saturating_sub(self.repeated_spill_buffer_bytes); + // Also publish attempted spill work and release inputs when the writer fails. + // Only batches actually produced for the writer contribute data-buffer bytes. + self.reservation.free(); self.metrics.memory_spilled_bytes.add(memory_spilled_bytes); self.pinned_buffers.clear(); - self.repeated_spill_buffer_bytes = 0; self.metrics.spill_count.add(1); write_result }) @@ -603,8 +567,7 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartition /// This function will slice input batch according to configured batch size and then /// shuffle rows into corresponding partition buffer. async fn insert_batch(&mut self, batch: RecordBatch) -> datafusion::common::Result<()> { - self.spill_accounted_input_buffers.clear(); - let result = with_trace_async("shuffle_insert_batch", self.tracing_enabled, || async { + with_trace_async("shuffle_insert_batch", self.tracing_enabled, || async { let start_time = Instant::now(); let mut start = 0; while start < batch.num_rows() { @@ -620,9 +583,7 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartition .add_duration(start_time.elapsed()); Ok(()) }) - .await; - self.spill_accounted_input_buffers.clear(); - result + .await } /// Writes buffered shuffled record batches into Arrow IPC bytes. @@ -674,6 +635,7 @@ mod tests { #[derive(Default)] struct FailingPartitionWriter { fail: bool, + consume_before_failure: bool, write_calls: usize, } @@ -689,6 +651,9 @@ mod tests { { self.write_calls += 1; if self.fail { + if self.consume_before_failure { + iter.next().transpose()?; + } return Err(DataFusionError::Execution( "injected write failure".to_string(), )); @@ -718,12 +683,19 @@ mod tests { #[tokio::test] async fn spill_write_error_releases_buffered_memory() { + check_spill_write_error_releases_buffered_memory(false).await; + check_spill_write_error_releases_buffered_memory(true).await; + } + + async fn check_spill_write_error_releases_buffered_memory(consume_before_failure: bool) { let batch = RecordBatch::try_from_iter([( "a", Arc::new(Int64Array::from(vec![0, 1, 2, 3])) as ArrayRef, )]) .unwrap(); - let buffer_bytes = batch.column(0).to_data().buffers()[0].capacity(); + let backing_buffer = batch.column(0).to_data().buffers()[0].clone(); + let buffer_bytes = backing_buffer.capacity(); + let input_owners = backing_buffer.strong_count(); let runtime = Arc::new(RuntimeEnv::default()); let metrics_set = ExecutionPlanMetricsSet::new(); let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( @@ -745,11 +717,10 @@ mod tests { assert_eq!(repartitioner.reservation.size(), 0); assert_eq!(runtime.memory_pool.reserved(), 0); assert!(repartitioner.pinned_buffers.is_empty()); - assert_eq!(repartitioner.spill_accounted_input_buffers.len(), 1); - assert_eq!(repartitioner.repeated_spill_buffer_bytes, 0); assert!(repartitioner.buffered_batches.is_empty()); assert!(repartitioner.partition_indices.iter().all(Vec::is_empty)); assert_eq!(repartitioner.partition_writer.write_calls, 2); + assert_eq!(backing_buffer.strong_count(), input_owners); let successful_spill_bytes = repartitioner.metrics.memory_spilled_bytes.value(); assert_eq!(repartitioner.spill_count(), 1); assert!(successful_spill_bytes > buffer_bytes); @@ -762,18 +733,23 @@ mod tests { .await .unwrap(); let reservation_before_failure = repartitioner.reservation.size(); - let repeated_before_failure = repartitioner.repeated_spill_buffer_bytes; + let index_bytes_before_failure = repartitioner + .partition_indices + .iter() + .map(|indices| indices.allocated_size()) + .sum::(); let metrics_before_failure = ( repartitioner.spill_count(), repartitioner.metrics.memory_spilled_bytes.value(), repartitioner.spilled_bytes(), repartitioner.data_size(), ); - assert_eq!(reservation_before_failure, successful_spill_bytes); + assert_eq!( + reservation_before_failure, + buffer_bytes + index_bytes_before_failure + ); assert_eq!(runtime.memory_pool.reserved(), reservation_before_failure); assert_eq!(repartitioner.pinned_buffers.len(), 1); - assert_eq!(repartitioner.spill_accounted_input_buffers.len(), 1); - assert_eq!(repeated_before_failure, buffer_bytes); assert_eq!(repartitioner.buffered_batches.len(), 1); assert_eq!( repartitioner @@ -790,7 +766,15 @@ mod tests { ); repartitioner.partition_writer.fail = true; - let error = repartitioner.spill(0).unwrap_err(); + repartitioner.partition_writer.consume_before_failure = consume_before_failure; + let materialized_bytes_before_failure = if consume_before_failure { + let first_output = + arrow::compute::interleave_record_batch(&[&batch], &[(0, 2)]).unwrap(); + get_record_batch_memory_size(&first_output) + } else { + 0 + }; + let error = repartitioner.spill().unwrap_err(); assert!(matches!( error, DataFusionError::Execution(message) if message == "injected write failure" @@ -798,11 +782,10 @@ mod tests { assert_eq!(repartitioner.reservation.size(), 0); assert_eq!(runtime.memory_pool.reserved(), 0); assert!(repartitioner.pinned_buffers.is_empty()); - assert_eq!(repartitioner.spill_accounted_input_buffers.len(), 1); - assert_eq!(repartitioner.repeated_spill_buffer_bytes, 0); assert!(repartitioner.buffered_batches.is_empty()); assert!(repartitioner.partition_indices.iter().all(Vec::is_empty)); assert_eq!(repartitioner.partition_writer.write_calls, 3); + assert_eq!(backing_buffer.strong_count(), input_owners); assert_eq!( ( repartitioner.spill_count(), @@ -812,10 +795,88 @@ mod tests { ), ( metrics_before_failure.0 + 1, - metrics_before_failure.1 + reservation_before_failure - repeated_before_failure, + metrics_before_failure.1 + + index_bytes_before_failure + + materialized_bytes_before_failure, metrics_before_failure.2, metrics_before_failure.3, ) ); } + + #[tokio::test] + 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) + .map(|i| format!("long string view payload {i}")) + .collect::>(); + let views = StringViewArray::from_iter( + strings + .iter() + .enumerate() + .map(|(i, value)| (i % 7 != 0).then_some(value.as_str())), + ); + let dictionary = DictionaryArray::::try_new( + Int32Array::from_iter((0..num_rows).map(|i| (i % 7 != 0).then_some((i % 3) as i32))), + Arc::new(StringArray::from(vec![ + "first dictionary value", + "second dictionary value", + "third dictionary value", + ])), + ) + .unwrap(); + let lists = ListArray::from_iter_primitive::( + (0..num_rows) + .map(|i| (i % 5 != 0).then_some(vec![Some(i as i64), None, Some(i as i64 + 1)])), + ); + let batch = RecordBatch::try_from_iter([ + ( + "a", + Arc::new(Int64Array::from_iter_values(0..num_rows as i64)) as ArrayRef, + ), + ("view", Arc::new(views) as ArrayRef), + ("dictionary", Arc::new(dictionary) as ArrayRef), + ("list", Arc::new(lists) as ArrayRef), + ]) + .unwrap(); + + // Views and dictionaries may keep shared backing allocations in the spill output. + // Compare that output footprint, not a globally unique count of input allocations. + let mut spill_bytes = Vec::new(); + for input_batch_rows in [num_rows, batch_size] { + let runtime = Arc::new(RuntimeEnv::default()); + let metrics_set = ExecutionPlanMetricsSet::new(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + FailingPartitionWriter::default(), + CometPartitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2), + ShufflePartitionerMetrics::new(&metrics_set, 0), + Arc::clone(&runtime), + batch_size, + false, + Some(1), + ) + .unwrap(); + for start in (0..num_rows).step_by(input_batch_rows) { + repartitioner + .insert_batch(batch.slice(start, input_batch_rows)) + .await + .unwrap(); + } + assert_eq!(repartitioner.spill_count(), num_rows / batch_size); + assert_eq!(repartitioner.reservation.size(), 0); + assert_eq!(runtime.memory_pool.reserved(), 0); + assert!(repartitioner.pinned_buffers.is_empty()); + assert!(repartitioner.buffered_batches.is_empty()); + assert!(repartitioner.partition_indices.iter().all(Vec::is_empty)); + spill_bytes.push(repartitioner.metrics.memory_spilled_bytes.value()); + } + assert!(spill_bytes[0] > 0); + assert_eq!(spill_bytes[0], spill_bytes[1]); + } } diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 21d71dc1221..097231320d5 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -564,7 +564,7 @@ mod test { assert!(!spill_writers[1].has_spill_file()); } - repartitioner.spill(0).unwrap(); + repartitioner.spill().unwrap(); // after spill, there should be spill files { @@ -640,13 +640,14 @@ mod test { ); } - /// Spill every slice of one shared Arrow allocation and verify that memory accounting - /// includes that backing allocation once per outer input batch, plus each slice's indices. + /// Spill the same rows with different input batching and allocation-sharing patterns. async fn shared_buffer_memory_spilled_bytes( max_buffer_bytes: Option, memory_limit: usize, input_batches: usize, - ) -> usize { + input_batch_rows: usize, + fresh_allocations: bool, + ) -> (usize, Vec) { let num_rows = 16_384usize; let batch_size = 1024usize; let num_partitions = 2; @@ -656,7 +657,7 @@ mod test { vec![Arc::new(Int64Array::from_iter_values(0..num_rows as i64))], ) .unwrap(); - let buffer_bytes = backing.get_array_memory_size(); + let value_bytes = num_rows * std::mem::size_of::(); let runtime_env = create_runtime(memory_limit); let metrics_set = ExecutionPlanMetricsSet::new(); @@ -689,7 +690,21 @@ mod test { .unwrap(); for _ in 0..input_batches { - repartitioner.insert_batch(backing.clone()).await.unwrap(); + for start in (0..num_rows).step_by(input_batch_rows) { + let end = (start + input_batch_rows).min(num_rows); + let input = if fresh_allocations { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from_iter_values( + start as i64..end as i64, + ))], + ) + .unwrap() + } else { + backing.slice(start, end - start) + }; + repartitioner.insert_batch(input).await.unwrap(); + } } repartitioner.shuffle_write().unwrap(); @@ -701,30 +716,32 @@ mod test { ); let spilled = memory_spilled_bytes.value(); - let minimum_backing_bytes = input_batches * buffer_bytes; + let minimum_data_bytes = input_batches * value_bytes; assert!( - spilled > minimum_backing_bytes, + spilled > minimum_data_bytes, "partition-index allocations must remain in memory spill accounting: \ - {spilled} bytes reported for {minimum_backing_bytes} backing bytes" + {spilled} bytes reported for {minimum_data_bytes} value bytes" ); // Each row receives one (batch index, row index) entry. Allow twice the logical index // size for Vec capacity rounding while still rejecting one full backing charge per slice. let maximum_index_bytes = input_batches * num_rows * std::mem::size_of::<(u32, u32)>() * 2; - let maximum_spilled = minimum_backing_bytes + maximum_index_bytes; + let maximum_spilled = minimum_data_bytes + maximum_index_bytes; assert!( spilled <= maximum_spilled, - "shared backing buffers must be charged once per input batch: \ + "spill size must reflect the rows and indices, not the full input backing per slice: \ {spilled} bytes reported, expected at most {maximum_spilled}" ); - spilled + (spilled, std::fs::read(dir.path().join("data.out")).unwrap()) } #[tokio::test] #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` - async fn max_buffer_spills_charge_shared_backing_once_per_input_batch() { - let one_batch = shared_buffer_memory_spilled_bytes(Some(8 * 1024), 512 * 1024, 1).await; - let two_batches = shared_buffer_memory_spilled_bytes(Some(8 * 1024), 512 * 1024, 2).await; + async fn max_buffer_spill_metrics_are_cumulative() { + let (one_batch, _) = + shared_buffer_memory_spilled_bytes(Some(8 * 1024), 512 * 1024, 1, 16_384, false).await; + let (two_batches, _) = + shared_buffer_memory_spilled_bytes(Some(8 * 1024), 512 * 1024, 2, 16_384, false).await; assert_eq!( two_batches, one_batch * 2, @@ -735,8 +752,43 @@ mod test { #[tokio::test] #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` - async fn rejected_reservations_charge_shared_backing_once_per_input_batch() { - shared_buffer_memory_spilled_bytes(None, 1, 1).await; + async fn rejected_reservations_count_materialized_spill_batches() { + shared_buffer_memory_spilled_bytes(None, 1, 1, 16_384, false).await; + } + + #[tokio::test] + #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` + async fn shared_buffer_spill_metrics_do_not_depend_on_input_batching() { + // Partial HashAggregate can emit the sixteen zero-copy slices separately. They must + // report the same memory spill size and write the same data as slicing one input here. + for (max_buffer_bytes, memory_limit) in [(Some(8 * 1024), 512 * 1024), (None, 1)] { + let (whole_bytes, whole_output) = shared_buffer_memory_spilled_bytes( + max_buffer_bytes, + memory_limit, + 1, + 16_384, + false, + ) + .await; + let (sliced_bytes, sliced_output) = + shared_buffer_memory_spilled_bytes(max_buffer_bytes, memory_limit, 1, 1024, false) + .await; + assert_eq!(sliced_bytes, whole_bytes); + assert_eq!(sliced_output, whole_output); + + // Independently allocated chunks must have the same spill representation as + // the shared zero-copy slices, without relying on allocation identities. + let (fresh_bytes, fresh_output) = + shared_buffer_memory_spilled_bytes(max_buffer_bytes, memory_limit, 1, 1024, true) + .await; + assert_eq!(fresh_bytes, whole_bytes); + assert_eq!(fresh_output, whole_output); + + let (repeated_bytes, _) = + shared_buffer_memory_spilled_bytes(max_buffer_bytes, memory_limit, 2, 1024, false) + .await; + assert_eq!(repeated_bytes, whole_bytes * 2); + } } /// Buffer `num_batches` batches through a `MultiPartitionShuffleRepartitioner` and return its From 9e593d18799bafa4b920843204629c806373e0ff Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Tue, 1 Sep 2026 23:54:06 +0000 Subject: [PATCH 2/2] fix: preserve per-spill input memory accounting --- docs/source/user-guide/latest/metrics.md | 8 + native/shuffle/src/metrics.rs | 5 +- .../src/partitioners/multi_partition.rs | 210 ++++++++++++++---- native/shuffle/src/shuffle_writer.rs | 35 ++- 4 files changed, 199 insertions(+), 59 deletions(-) diff --git a/docs/source/user-guide/latest/metrics.md b/docs/source/user-guide/latest/metrics.md index a8b2d6bdb8a..f1804e5c13c 100644 --- a/docs/source/user-guide/latest/metrics.md +++ b/docs/source/user-guide/latest/metrics.md @@ -50,6 +50,14 @@ enabled and uncompressed when `spark.shuffle.compress=false`. Memory spill bytes partition-index data rather than their on-disk size. These values also appear in Spark's task metrics and Spark UI as `diskBytesSpilled` and `memoryBytesSpilled`, respectively. +Memory spill bytes are cumulative across spills, not a peak-memory measurement or a count of +allocations unique across the whole task. Each spill counts the full capacity of its buffered +input allocations, deduplicating buffers shared by columns or batches in that spill, plus its +partition-index allocations. If a later spill buffers the same backing allocation again, it +contributes again. Whether input slices arrive in one batch or separate batches does not change +the accounting for identical spill boundaries. Other operators may still own the same buffers, +so this measures memory released from shuffle buffering, not necessarily a drop in process memory. + ## Native Metrics Setting `spark.comet.explain.native.enabled=true` will cause native plans to be logged in each executor. Metrics are diff --git a/native/shuffle/src/metrics.rs b/native/shuffle/src/metrics.rs index 5648f8649f1..855bb950111 100644 --- a/native/shuffle/src/metrics.rs +++ b/native/shuffle/src/metrics.rs @@ -46,9 +46,8 @@ pub(crate) struct ShufflePartitionerMetrics { /// total spilled bytes during the execution of the operator pub(crate) spilled_bytes: Count, - /// Total buffer size of materialized spill batches before compression, plus the - /// partition-index allocations released by spills. Measured from spill output rather - /// than input batch boundaries; not a count of globally unique input allocations. + /// Cumulative input backing-buffer and partition-index capacity released by spills. + /// Shared input allocations are counted once per spill, not once per input batch. pub(crate) memory_spilled_bytes: Count, /// The original size of spilled data. Different to `spilled_bytes` because of compression. diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index ab2e4dee1cd..f7a8d620de7 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -21,7 +21,6 @@ use crate::partitioners::ShufflePartitioner; use crate::writers::PartitionWriter; use crate::{comet_partitioning, CometPartitioning}; use arrow::array::{Array, ArrayData, ArrayRef, RecordBatch}; -use datafusion::common::utils::memory::get_record_batch_memory_size; use datafusion::common::utils::proxy::VecAllocExt; use datafusion::common::{DataFusionError, HashSet}; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; @@ -470,7 +469,7 @@ impl MultiPartitionShuffleRepartitioner { .max_buffer_bytes .is_some_and(|limit| self.reservation.size() >= limit) { - self.spill()?; + self.spill(if reservation_failed { mem_growth } else { 0 })?; } Ok(()) @@ -506,7 +505,7 @@ impl MultiPartitionShuffleRepartitioner { PartitionedBatchesProducer::new(buffered_batches, indices, self.batch_size) } - pub(crate) fn spill(&mut self) -> datafusion::common::Result<()> { + pub(crate) fn spill(&mut self, unreserved_bytes: usize) -> datafusion::common::Result<()> { log::info!( "ShuffleRepartitioner spilling {} bytes to its partition writer ({} previous spills)", self.used(), @@ -520,34 +519,23 @@ impl MultiPartitionShuffleRepartitioner { with_trace("shuffle_spill", self.tracing_enabled, || { let num_output_partitions = self.partition_indices.len(); - // The reservation measures pinned input allocations, not the size of the data - // being spilled: many input batches may be slices of the same large allocation. - // Measure the already-materialized output batches instead, before compression, - // and include the partition-index allocations released by this spill. This is - // cumulative across spills without retaining input buffers or their addresses. - let mut memory_spilled_bytes = self - .partition_indices - .iter() - .map(|indices| indices.allocated_size()) - .sum::(); let write_result = { let mut partitioned_batches = self.partitioned_batches(); (0..num_output_partitions).try_for_each(|partition_id| { - let mut batches = partitioned_batches - .produce(partition_id, &self.metrics.interleave_time) - .inspect(|result| { - if let Ok(batch) = result { - memory_spilled_bytes += get_record_batch_memory_size(batch); - } - }); - self.partition_writer - .write(partition_id, &mut batches, &self.metrics) + self.partition_writer.write( + partition_id, + &mut partitioned_batches + .produce(partition_id, &self.metrics.interleave_time), + &self.metrics, + ) }) }; - // Also publish attempted spill work and release inputs when the writer fails. - // Only batches actually produced for the writer contribute data-buffer bytes. - self.reservation.free(); + // 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); self.metrics.memory_spilled_bytes.add(memory_spilled_bytes); self.pinned_buffers.clear(); self.metrics.spill_count.add(1); @@ -681,6 +669,97 @@ mod tests { } } + async fn check_spill_metrics_count_input_buffers(batch: RecordBatch, input_bytes: usize) { + let runtime = Arc::new(RuntimeEnv::default()); + let metrics_set = ExecutionPlanMetricsSet::new(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + FailingPartitionWriter::default(), + CometPartitioning::RoundRobin(2, 1), + ShufflePartitionerMetrics::new(&metrics_set, 0), + Arc::clone(&runtime), + 64, + false, + None, + ) + .unwrap(); + repartitioner.insert_batch(batch).await.unwrap(); + let index_bytes = repartitioner + .partition_indices + .iter() + .map(|indices| indices.allocated_size()) + .sum::(); + let reserved_bytes = repartitioner.reservation.size(); + assert_eq!(repartitioner.spill_count(), 0); + assert_eq!(reserved_bytes, input_bytes + index_bytes); + + repartitioner.spill(0).unwrap(); + + assert_eq!( + repartitioner.metrics.memory_spilled_bytes.value(), + reserved_bytes + ); + assert_eq!(runtime.memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn spill_metrics_count_input_allocation_capacity() { + // Only 8 KiB of rows are populated, but the input pins the full 1 MiB allocation. + let mut values = Vec::with_capacity(1024 * 1024 / std::mem::size_of::()); + values.extend(0..1024i64); + let values = Int64Array::from(values); + let input_bytes = values.get_buffer_memory_size(); + assert_eq!(input_bytes, 1024 * 1024); + let batch = RecordBatch::try_from_iter([("a", Arc::new(values) as ArrayRef)]).unwrap(); + check_spill_metrics_count_input_buffers(batch, input_bytes).await; + } + + #[tokio::test] + async fn spill_metrics_count_aliased_input_columns_once() { + let values: ArrayRef = Arc::new(Int64Array::from_iter_values(0..1024)); + let input_bytes = values.get_buffer_memory_size(); + let batch = + RecordBatch::try_from_iter([("a", Arc::clone(&values)), ("b", values)]).unwrap(); + check_spill_metrics_count_input_buffers(batch, input_bytes).await; + } + + #[tokio::test] + async fn spill_metrics_count_shared_view_payload_once() { + use arrow::array::StringViewArray; + + // Interleaved output batches retain the same out-of-line input payload buffers. + let values = StringViewArray::from_iter_values( + (0..1024).map(|i| format!("shared string view payload {i}")), + ); + // Partition on integers, since the native hasher does not support string views. + let keys = Int64Array::from_iter_values(0..1024); + let input_bytes = keys.get_buffer_memory_size() + values.get_buffer_memory_size(); + let batch = RecordBatch::try_from_iter([ + ("key", Arc::new(keys) as ArrayRef), + ("a", Arc::new(values) as ArrayRef), + ]) + .unwrap(); + check_spill_metrics_count_input_buffers(batch, input_bytes).await; + } + + #[tokio::test] + async fn spill_metrics_count_shared_dictionary_values_once() { + use arrow::array::{DictionaryArray, Int32Array, StringArray}; + use arrow::datatypes::Int32Type; + + let values = DictionaryArray::::try_new( + Int32Array::from_iter_values((0..1024).map(|i| i % 2)), + Arc::new(StringArray::from(vec![ + "first shared dictionary value", + "second shared dictionary value", + ])), + ) + .unwrap(); + let input_bytes = values.get_buffer_memory_size(); + let batch = RecordBatch::try_from_iter([("a", Arc::new(values) as ArrayRef)]).unwrap(); + check_spill_metrics_count_input_buffers(batch, input_bytes).await; + } + #[tokio::test] async fn spill_write_error_releases_buffered_memory() { check_spill_write_error_releases_buffered_memory(false).await; @@ -767,14 +846,7 @@ mod tests { repartitioner.partition_writer.fail = true; repartitioner.partition_writer.consume_before_failure = consume_before_failure; - let materialized_bytes_before_failure = if consume_before_failure { - let first_output = - arrow::compute::interleave_record_batch(&[&batch], &[(0, 2)]).unwrap(); - get_record_batch_memory_size(&first_output) - } else { - 0 - }; - let error = repartitioner.spill().unwrap_err(); + let error = repartitioner.spill(0).unwrap_err(); assert!(matches!( error, DataFusionError::Execution(message) if message == "injected write failure" @@ -795,15 +867,77 @@ mod tests { ), ( metrics_before_failure.0 + 1, - metrics_before_failure.1 - + index_bytes_before_failure - + materialized_bytes_before_failure, + metrics_before_failure.1 + reservation_before_failure, metrics_before_failure.2, metrics_before_failure.3, ) ); } + #[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([( + "a", + Arc::new(Int64Array::from_iter_values(0..64)) as ArrayRef, + )]) + .unwrap(); + let next_batch = if share_buffer { + batch.clone() + } else { + RecordBatch::try_from_iter([( + "a", + Arc::new(Int64Array::from_iter_values(64..128)) as ArrayRef, + )]) + .unwrap() + }; + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_limit(1024, 1.0) + .build() + .unwrap(), + ); + let metrics_set = ExecutionPlanMetricsSet::new(); + let mut repartitioner = MultiPartitionShuffleRepartitioner::try_new( + 0, + FailingPartitionWriter::default(), + CometPartitioning::RoundRobin(2, 0), + ShufflePartitionerMetrics::new(&metrics_set, 0), + Arc::clone(&runtime), + 64, + false, + None, + ) + .unwrap(); + let row_indices = (0..64).collect::>(); + repartitioner + .buffer_partitioned_batch_may_spill(batch, &row_indices, &[0, 32, 64]) + .await + .unwrap(); + // The first reservation includes 512 bytes of input and 512 bytes of indices. + assert_eq!(repartitioner.reservation.size(), 1024); + assert_eq!(repartitioner.spill_count(), 0); + + repartitioner.partition_writer.fail = fail; + repartitioner.partition_writer.consume_before_failure = consume_before_failure; + let result = repartitioner + .buffer_partitioned_batch_may_spill(next_batch, &row_indices, &[0, 32, 64]) + .await; + assert_eq!(result.is_err(), fail); + // The indices grow by 512 bytes. Only independent input adds another 512. + assert_eq!( + repartitioner.metrics.memory_spilled_bytes.value(), + if share_buffer { 1536 } else { 2048 } + ); + assert_eq!(repartitioner.spill_count(), 1); + assert_eq!(runtime.memory_pool.reserved(), 0); + } + } + } + #[tokio::test] async fn heterogeneous_spill_metrics_do_not_depend_on_input_batching() { use arrow::array::{DictionaryArray, Int32Array, ListArray, StringArray, StringViewArray}; @@ -845,8 +979,8 @@ mod tests { ]) .unwrap(); - // Views and dictionaries may keep shared backing allocations in the spill output. - // Compare that output footprint, not a globally unique count of input allocations. + // Shared child allocations contribute once per spill, independently of whether the + // caller or insert_batch slices the input. They are not counted per output batch. let mut spill_bytes = Vec::new(); for input_batch_rows in [num_rows, batch_size] { let runtime = Arc::new(RuntimeEnv::default()); diff --git a/native/shuffle/src/shuffle_writer.rs b/native/shuffle/src/shuffle_writer.rs index 097231320d5..badcede116c 100644 --- a/native/shuffle/src/shuffle_writer.rs +++ b/native/shuffle/src/shuffle_writer.rs @@ -564,7 +564,7 @@ mod test { assert!(!spill_writers[1].has_spill_file()); } - repartitioner.spill().unwrap(); + repartitioner.spill(0).unwrap(); // after spill, there should be spill files { @@ -657,7 +657,6 @@ mod test { vec![Arc::new(Int64Array::from_iter_values(0..num_rows as i64))], ) .unwrap(); - let value_bytes = num_rows * std::mem::size_of::(); let runtime_env = create_runtime(memory_limit); let metrics_set = ExecutionPlanMetricsSet::new(); @@ -689,6 +688,7 @@ mod test { ) .unwrap(); + let mut expected_buffer_bytes = 0; for _ in 0..input_batches { for start in (0..num_rows).step_by(input_batch_rows) { let end = (start + input_batch_rows).min(num_rows); @@ -703,6 +703,9 @@ mod test { } else { backing.slice(start, end - start) }; + // Each internal chunk spills and releases its full pinned input allocation. + expected_buffer_bytes += input.column(0).to_data().buffers()[0].capacity() + * input.num_rows().div_ceil(batch_size); repartitioner.insert_batch(input).await.unwrap(); } } @@ -716,20 +719,16 @@ mod test { ); let spilled = memory_spilled_bytes.value(); - let minimum_data_bytes = input_batches * value_bytes; - assert!( - spilled > minimum_data_bytes, - "partition-index allocations must remain in memory spill accounting: \ - {spilled} bytes reported for {minimum_data_bytes} value bytes" - ); // Each row receives one (batch index, row index) entry. Allow twice the logical index - // size for Vec capacity rounding while still rejecting one full backing charge per slice. - let maximum_index_bytes = input_batches * num_rows * std::mem::size_of::<(u32, u32)>() * 2; - let maximum_spilled = minimum_data_bytes + maximum_index_bytes; + // size for Vec capacity rounding. Buffer capacity contributes again on every spill, + // regardless of whether the caller or insert_batch sliced the input. + let minimum_index_bytes = input_batches * num_rows * std::mem::size_of::<(u32, u32)>(); assert!( - spilled <= maximum_spilled, - "spill size must reflect the rows and indices, not the full input backing per slice: \ - {spilled} bytes reported, expected at most {maximum_spilled}" + (expected_buffer_bytes + minimum_index_bytes + ..=expected_buffer_bytes + minimum_index_bytes * 2) + .contains(&spilled), + "each spill must count its pinned input buffers and indices: \ + {spilled} bytes reported for {expected_buffer_bytes} buffer bytes" ); (spilled, std::fs::read(dir.path().join("data.out")).unwrap()) @@ -752,7 +751,7 @@ mod test { #[tokio::test] #[cfg_attr(miri, ignore)] // miri can't call foreign function `ZSTD_createCCtx` - async fn rejected_reservations_count_materialized_spill_batches() { + async fn rejected_reservations_count_pinned_input_buffers() { shared_buffer_memory_spilled_bytes(None, 1, 1, 16_384, false).await; } @@ -776,12 +775,12 @@ mod test { assert_eq!(sliced_bytes, whole_bytes); assert_eq!(sliced_output, whole_output); - // Independently allocated chunks must have the same spill representation as - // the shared zero-copy slices, without relying on allocation identities. + // Independently allocated chunks pin less memory per spill, even though their + // serialized output is identical to the shared zero-copy slices. let (fresh_bytes, fresh_output) = shared_buffer_memory_spilled_bytes(max_buffer_bytes, memory_limit, 1, 1024, true) .await; - assert_eq!(fresh_bytes, whole_bytes); + assert!(fresh_bytes < whole_bytes); assert_eq!(fresh_output, whole_output); let (repeated_bytes, _) =