You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/18 11:46:56 UTC
(datafusion) branch gh-readonly-queue/main/pr-26366-0ee39784e8702c0583a73f1a88fcae8356897f56 created (now 3d9037cb0f)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-26366-0ee39784e8702c0583a73f1a88fcae8356897f56
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 3d9037cb0f fix: avoid quadratic row-buffer accounting in sort merges (#26366)
This branch includes the following new commits:
new 3d9037cb0f fix: avoid quadratic row-buffer accounting in sort merges (#26366)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: avoid quadratic row-buffer accounting in sort merges (#26366)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26366-0ee39784e8702c0583a73f1a88fcae8356897f56
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 3d9037cb0fdc4965cd4fd5ea1f5e6f5ff7927856
Author: Geoffrey Claude <ge...@gmail.com>
AuthorDate: Sun Oct 18 11:46:32 2026 +0000
fix: avoid quadratic row-buffer accounting in sort merges (#26366)
## Which issue does this PR close?
Follow-up to #25372, related to #23174.
## Rationale for this change
Multi-key sorts over many small batches spend quadratic time accounting
for retained row buffers, delaying their first merged output.
The merge caches one encoded row buffer per input stream. Each new batch
currently scans every cache slot to recalculate the reservation. Reading
the first batch from each of N streams therefore requires N full scans.
A large merge can spend seconds doing this bookkeeping in a single Tokio
poll.
## What changes are included in this PR?
- Track each stream's reserved bytes and update the total by that
stream's change, making each accounting update constant-time.
- Keep the existing charge while a buffer is reused. Failed buffer
growth and stream completion release only that stream's charge.
- Add a SQL-based case to the existing `sort` benchmark that sorts two
million rows loaded in 32-row batches. Two sort keys exercise the
affected row-format merge. Data generation happens before timing.
## What is the testing strategy for this PR?
A focused regression test fills a bounded memory pool, then rejects
growth of a reused row buffer. It checks that the failed stream releases
its previous reservation while another stream remains accounted for.
The SQL benchmark runs a single query, with no concurrent workload:
```bash
cargo bench -p datafusion --bench sort -- small_batches/two_keys
```
[Benchmark bot
results](https://github.com/apache/datafusion/pull/26366#issuecomment-6194773657):
**13.9 s → 3.4 s (4.09× faster)** for `small_batches/two_keys`.
## Are there any user-facing changes?
Large row-format merges perform less bookkeeping while preserving their
memory reservations. The internal reservation invariant is documented
alongside the per-stream sizes.
Developed with AI assistance.
---
datafusion/core/benches/sort.rs | 61 +++++++++++++++++++++++++++-
datafusion/physical-plan/src/sorts/stream.rs | 56 +++++++++++++++++++++----
2 files changed, 108 insertions(+), 9 deletions(-)
diff --git a/datafusion/core/benches/sort.rs b/datafusion/core/benches/sort.rs
index ac4be5b8b2..5dcdfb3eac 100644
--- a/datafusion/core/benches/sort.rs
+++ b/datafusion/core/benches/sort.rs
@@ -1065,5 +1065,64 @@ fn with_extra_columns(batches: Vec<RecordBatch>, n: usize) -> Vec<RecordBatch> {
.collect()
}
-criterion_group!(benches, criterion_benchmark, sort_axis_benchmark);
+fn small_batches_benchmark(c: &mut Criterion) {
+ let mut group = c.benchmark_group("small_batches");
+ group.sample_size(10);
+ group.sampling_mode(criterion::SamplingMode::Flat);
+ group.bench_function("two_keys", |b| {
+ let runtime = Runtime::new().unwrap();
+ let ctx = SessionContext::new_with_config(
+ SessionConfig::new()
+ .with_target_partitions(1)
+ .with_batch_size(32),
+ );
+
+ // Store small input batches before timing, then restore normal output batches.
+ runtime.block_on(async {
+ ctx.sql(
+ "CREATE TABLE sort_input AS \
+ SELECT value % 997 AS sort_key, value FROM range(2000000)",
+ )
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ ctx.sql("SET datafusion.execution.batch_size = 8192")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ });
+
+ // Two keys exercise row-format merging of the small sorted runs.
+ b.iter(|| {
+ runtime.block_on(async {
+ let mut stream = ctx
+ .sql(
+ "SELECT sort_key, value FROM sort_input ORDER BY sort_key, value",
+ )
+ .await
+ .unwrap()
+ .execute_stream()
+ .await
+ .unwrap();
+ let mut row_count = 0;
+ while let Some(batch) = stream.next().await {
+ row_count += batch.unwrap().num_rows();
+ }
+ row_count
+ })
+ });
+ });
+ group.finish();
+}
+
+criterion_group!(
+ benches,
+ criterion_benchmark,
+ sort_axis_benchmark,
+ small_batches_benchmark
+);
criterion_main!(benches);
diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs
index d7bc628225..b1708beb51 100644
--- a/datafusion/physical-plan/src/sorts/stream.rs
+++ b/datafusion/physical-plan/src/sorts/stream.rs
@@ -103,6 +103,9 @@ impl FusedStreams {
#[derive(Debug)]
struct ReusableRows {
inner: Vec<Option<Arc<Rows>>>,
+ /// Each stream's charge stays reserved while its buffer is taken for reuse.
+ /// The sum of these sizes equals `reservation.size()`.
+ reserved_sizes: Vec<usize>,
reservation: MemoryReservation,
}
@@ -130,23 +133,27 @@ impl ReusableRows {
/// reservation is free for the same reason. A pool that cannot cover the buffer
/// fails the query here, as it did before the buffer was cached at all.
fn save(&mut self, stream_idx: usize, rows: &Arc<Rows>) -> Result<()> {
- self.inner[stream_idx] = Some(Arc::clone(rows));
- let retained = self.retained_size();
+ let size = rows.size();
+ let retained = self.reservation.size() - self.reserved_sizes[stream_idx] + size;
debug_assert!(retained >= self.reservation.size());
if let Err(e) = self.reservation.try_resize(retained) {
- self.inner[stream_idx] = None;
- self.reservation.try_resize(self.retained_size())?;
+ self.release(stream_idx);
return Err(e);
}
+ self.inner[stream_idx] = Some(Arc::clone(rows));
+ self.reserved_sizes[stream_idx] = size;
Ok(())
}
// drop whatever a finished stream was holding
fn release(&mut self, stream_idx: usize) {
debug_assert!(self.reservation.size() >= self.retained_size());
- if let Some(rows) = self.inner[stream_idx].take() {
- self.reservation.shrink(rows.size());
- }
+ self.inner[stream_idx] = None;
+ let reserved = mem::take(&mut self.reserved_sizes[stream_idx]);
+ self.reservation.resize(self.reservation.size() - reserved);
+ // Verify the incremental accounting against a full recomputation, avoiding
+ // the cost of scanning all buffers in release builds.
+ debug_assert_eq!(self.reservation.size(), self.retained_size());
}
fn retained_size(&self) -> usize {
@@ -200,6 +207,7 @@ impl RowCursorStream {
streams: FusedStreams(streams),
rows: ReusableRows {
inner: vec![None; stream_count],
+ reserved_sizes: vec![0; stream_count],
reservation: rows_reservation,
},
})
@@ -442,7 +450,7 @@ mod tests {
use datafusion_common::DataFusionError;
use datafusion_execution::RecordBatchStream;
use datafusion_execution::memory_pool::{
- GreedyMemoryPool, MemoryConsumer, MemoryPool,
+ GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool,
};
use datafusion_physical_expr::expressions::col;
use futures::Stream;
@@ -907,4 +915,36 @@ mod tests {
Ok(())
}
+
+ // The buffer has already been taken out of the cache when its reservation fails
+ // to grow. Cleanup must still release its previous charge.
+ #[test]
+ fn failed_buffer_growth_releases_only_its_stream_reservation() -> Result<()> {
+ let small = (64, 8);
+ let large = (2048, 128);
+ let mut harness =
+ RowCursorHarness::with_shapes(vec![vec![small, large], vec![small]])?;
+
+ // Cache stream 1's buffer and record the reservation that must survive.
+ assert!(harness.poll_and_drop_cursor(1)?);
+ let expected_remaining = harness.reserved();
+ assert!(harness.poll_and_drop_cursor(0)?);
+
+ // Leave no space for stream 0's larger batch.
+ let MemoryLimit::Finite(limit) = harness.pool.memory_limit() else {
+ unreachable!("the harness uses a bounded memory pool");
+ };
+ let competitor = MemoryConsumer::new("competitor").register(&harness.pool);
+ competitor.try_grow(limit - harness.reserved())?;
+
+ assert!(matches!(
+ harness.poll_cursor(0),
+ Err(DataFusionError::ResourcesExhausted(_))
+ ));
+ drop(competitor);
+
+ // Only stream 1's buffer and the shared converter remain reserved.
+ assert_eq!(harness.reserved(), expected_remaining);
+ Ok(())
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org