You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/04/30 22:26:29 UTC
(datafusion-ballista) branch main updated: perf(sort-shuffle): fix performance regression caused by datafusion upgrade (#1626)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new de8034dc0 perf(sort-shuffle): fix performance regression caused by datafusion upgrade (#1626)
de8034dc0 is described below
commit de8034dc0730fdcb488d79b3d61db3f78d0f214c
Author: Andy Grove <ag...@apache.org>
AuthorDate: Thu Apr 30 16:26:25 2026 -0600
perf(sort-shuffle): fix performance regression caused by datafusion upgrade (#1626)
---
.../sort_shuffle/partitioned_batch_iterator.rs | 36 +++++++++++++++-
.../core/src/execution_plans/sort_shuffle/spill.rs | 26 ++++++-----
.../src/execution_plans/sort_shuffle/writer.rs | 50 +++++++++++++++++-----
benchmarks/src/bin/shuffle_bench.rs | 3 ++
4 files changed, 94 insertions(+), 21 deletions(-)
diff --git a/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs b/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs
index 896cea8ef..2b92f414a 100644
--- a/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs
+++ b/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs
@@ -18,10 +18,44 @@
//! Iterator that materializes per-partition rows into well-sized
//! `RecordBatch`es using `arrow::compute::interleave_record_batch`.
+use datafusion::arrow::array::{Array, ArrayRef, BinaryViewArray, StringViewArray};
use datafusion::arrow::compute::interleave_record_batch;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::common::DataFusionError;
use datafusion::error::Result;
+use std::sync::Arc;
+
+/// Compacts `Utf8View` / `BinaryView` columns by running `gc()` on them.
+///
+/// `interleave_record_batch` preserves the original data buffers for view
+/// arrays, so the output references every source data buffer the inputs
+/// touched even if it only emits a few rows. Without compaction, downstream
+/// IPC writes serialize all those source buffers (potentially hundreds of MB)
+/// rather than just the bytes the views point at.
+fn compact_view_columns(batch: RecordBatch) -> Result<RecordBatch> {
+ let mut changed = false;
+ let columns: Vec<ArrayRef> = batch
+ .columns()
+ .iter()
+ .map(|c| -> ArrayRef {
+ if let Some(a) = c.as_any().downcast_ref::<StringViewArray>() {
+ changed = true;
+ Arc::new(a.gc())
+ } else if let Some(a) = c.as_any().downcast_ref::<BinaryViewArray>() {
+ changed = true;
+ Arc::new(a.gc())
+ } else {
+ c.clone()
+ }
+ })
+ .collect();
+ if !changed {
+ return Ok(batch);
+ }
+ RecordBatch::try_new(batch.schema(), columns).map_err(|e| {
+ DataFusionError::ArrowError(Box::new(e), Some(DataFusionError::get_back_trace()))
+ })
+}
/// Iterator over per-partition output `RecordBatch`es.
///
@@ -79,7 +113,7 @@ impl<'a> Iterator for PartitionedBatchIterator<'a> {
self.pos = end;
match interleave_record_batch(&self.batch_refs, &self.scratch) {
- Ok(batch) => Some(Ok(batch)),
+ Ok(batch) => Some(compact_view_columns(batch)),
Err(e) => Some(Err(DataFusionError::ArrowError(
Box::new(e),
Some(DataFusionError::get_back_trace()),
diff --git a/ballista/core/src/execution_plans/sort_shuffle/spill.rs b/ballista/core/src/execution_plans/sort_shuffle/spill.rs
index 0f7cf8947..4dbb63146 100644
--- a/ballista/core/src/execution_plans/sort_shuffle/spill.rs
+++ b/ballista/core/src/execution_plans/sort_shuffle/spill.rs
@@ -51,8 +51,10 @@ pub struct SpillManager {
active_writers: HashMap<usize, StreamWriter<BufWriter<File>>>,
/// Compression codec for spill files
compression: CompressionType,
- /// Total number of spills performed (counts batches, not events)
- total_spills: usize,
+ /// Total number of batches written across all spill files. One call to
+ /// `spill_all_partitions` typically increments this multiple times (once
+ /// per partition that had buffered rows).
+ total_spilled_batches: u64,
/// Total bytes spilled
total_bytes_spilled: u64,
/// Per-partition counters: partition_id -> (batches, rows, bytes)
@@ -91,7 +93,7 @@ impl SpillManager {
spill_files: HashMap::new(),
active_writers: HashMap::new(),
compression,
- total_spills: 0,
+ total_spilled_batches: 0,
total_bytes_spilled: 0,
partition_counters: HashMap::new(),
})
@@ -138,7 +140,7 @@ impl SpillManager {
entry.1 += batch.num_rows() as u64;
entry.2 += bytes_written;
- self.total_spills += 1;
+ self.total_spilled_batches += 1;
self.total_bytes_spilled += bytes_written;
Ok(bytes_written)
@@ -186,9 +188,13 @@ impl SpillManager {
Ok(())
}
- /// Returns the total number of batches spilled across all partitions.
- pub fn total_spills(&self) -> usize {
- self.total_spills
+ /// Returns the total number of batches written to spill files across all
+ /// partitions. Note this counts batches, not spill *events*: a single
+ /// memory-pressure event in the writer typically produces one batch per
+ /// non-empty output partition. Spill-event accounting lives at the writer
+ /// layer because the spill manager only sees batch-level calls.
+ pub fn total_spilled_batches(&self) -> u64 {
+ self.total_spilled_batches
}
/// Returns the total bytes spilled to disk.
@@ -228,7 +234,7 @@ impl std::fmt::Debug for SpillManager {
.field("spill_dir", &self.spill_dir)
.field("spill_files", &self.spill_files)
.field("compression", &self.compression)
- .field("total_spills", &self.total_spills)
+ .field("total_spilled_batches", &self.total_spilled_batches)
.field("total_bytes_spilled", &self.total_bytes_spilled)
.finish()
}
@@ -281,7 +287,7 @@ mod tests {
assert!(manager.has_spill_files(0));
assert!(!manager.has_spill_files(1));
- assert_eq!(manager.total_spills(), 2);
+ assert_eq!(manager.total_spilled_batches(), 2);
manager.finish_writers()?;
@@ -314,7 +320,7 @@ mod tests {
assert!(manager.has_spill_files(0));
assert!(manager.has_spill_files(1));
- assert_eq!(manager.total_spills(), 2);
+ assert_eq!(manager.total_spilled_batches(), 2);
manager.finish_writers()?;
diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
index ea9479340..0d8ab0a16 100644
--- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs
+++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs
@@ -104,7 +104,10 @@ struct SortShuffleWriteMetrics {
input_rows: metrics::Count,
/// Number of output rows
output_rows: metrics::Count,
- /// Number of batches spilled to disk
+ /// Number of times the writer flushed buffered partitions to disk under
+ /// memory pressure. One event typically writes one batch per non-empty
+ /// output partition, so this is strictly less than or equal to the number
+ /// of batches written into spill files.
spill_count: metrics::Count,
/// Bytes spilled to disk
spill_bytes: metrics::Count,
@@ -236,6 +239,7 @@ impl SortShuffleWriterExec {
.register(&context.runtime_env().memory_pool);
let mut hash_buffer: Vec<u64> = Vec::new();
+ let mut spill_events: u64 = 0;
while let Some(result) = stream.next().await {
let input_batch = result?;
@@ -260,13 +264,29 @@ impl SortShuffleWriterExec {
if reservation.try_grow(growth).is_err() {
let spill_timer = metrics.spill_time.timer();
- spill_all_partitions(
+ let (event_batches, event_bytes) = spill_all_partitions(
&mut buffered,
&mut spill_manager,
&mut reservation,
config.batch_size,
)?;
spill_timer.done();
+
+ if event_batches > 0 {
+ spill_events += 1;
+ info!(
+ "Sort shuffle writer for input partition {} spilled \
+ event #{}: {} batches, {} bytes \
+ (cumulative: {} events, {} batches, {} bytes)",
+ input_partition,
+ spill_events,
+ event_batches,
+ event_bytes,
+ spill_events,
+ spill_manager.total_spilled_batches(),
+ spill_manager.total_bytes_spilled(),
+ );
+ }
}
}
@@ -288,7 +308,7 @@ impl SortShuffleWriterExec {
)?;
timer.done();
- metrics.spill_count.add(spill_manager.total_spills());
+ metrics.spill_count.add(spill_events as usize);
metrics
.spill_bytes
.add(spill_manager.total_bytes_spilled() as usize);
@@ -298,7 +318,7 @@ impl SortShuffleWriterExec {
// Snapshot spill counters before cleanup (cleanup doesn't touch them
// but we want to be explicit about ordering for the log line below).
- let total_spills = spill_manager.total_spills();
+ let total_spilled_batches = spill_manager.total_spilled_batches();
let total_bytes_spilled = spill_manager.total_bytes_spilled();
spill_manager
@@ -310,12 +330,14 @@ impl SortShuffleWriterExec {
info!(
"Sort shuffle write for partition {} completed in {} seconds. \
- Output: {:?}, Index: {:?}, Spill batches: {}, Spill bytes: {}",
+ Output: {:?}, Index: {:?}, Spill events: {}, Spill batches: {}, \
+ Spill bytes: {}",
input_partition,
now.elapsed().as_secs(),
data_path,
index_path,
- total_spills,
+ spill_events,
+ total_spilled_batches,
total_bytes_spilled
);
@@ -342,15 +364,21 @@ impl SortShuffleWriterExec {
/// indices through `PartitionedBatchIterator` and appends each yielded batch
/// to that partition's spill file. After this call, `buffered.is_empty()` is
/// true and `reservation.size() == 0`.
+///
+/// Returns `(batches_written, bytes_written)` for this single spill event so
+/// the caller can log per-event diagnostics. A return of `(0, 0)` means there
+/// was nothing buffered and no event occurred.
fn spill_all_partitions(
buffered: &mut BufferedBatches,
spill_manager: &mut SpillManager,
reservation: &mut MemoryReservation,
batch_size: usize,
-) -> Result<()> {
+) -> Result<(u64, u64)> {
if buffered.is_empty() {
- return Ok(());
+ return Ok((0, 0));
}
+ let mut batches_written: u64 = 0;
+ let mut bytes_written: u64 = 0;
let (batches, indices) = buffered.take();
for (partition_id, partition_indices) in indices.iter().enumerate() {
if partition_indices.is_empty() {
@@ -359,13 +387,15 @@ fn spill_all_partitions(
let iter = PartitionedBatchIterator::new(&batches, partition_indices, batch_size);
for result in iter {
let batch = result?;
- spill_manager
+ let written = spill_manager
.spill(partition_id, &batch)
.map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
+ batches_written += 1;
+ bytes_written += written;
}
}
reservation.free();
- Ok(())
+ Ok((batches_written, bytes_written))
}
/// Finalizes the output by writing the consolidated data file and index file.
diff --git a/benchmarks/src/bin/shuffle_bench.rs b/benchmarks/src/bin/shuffle_bench.rs
index 7b7038a65..d4f248b46 100644
--- a/benchmarks/src/bin/shuffle_bench.rs
+++ b/benchmarks/src/bin/shuffle_bench.rs
@@ -431,6 +431,9 @@ fn describe_schema(schema: &datafusion::arrow::datatypes::Schema) -> String {
}
fn main() {
+ env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
+ .init();
+
let args = Args::parse();
let writer_kind = parse_writer(&args.writer).unwrap_or_else(|e| {
eprintln!("error: {e}");
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org