You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "MisterRaindrop (via GitHub)" <gi...@apache.org> on 2026/10/13 13:44:14 UTC
[PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
MisterRaindrop opened a new pull request, #26264:
URL: https://github.com/apache/datafusion/pull/26264
## Which issue does this PR close?
- Part of #25804 (finding 9: "The final spill merge grows its reservation until the pool refuses").
## Rationale for this change
A sort that has spilled finishes with a merge of its spilled runs. That merge
seats as many runs as the memory pool will admit (the fan-in is unlimited by
default), and it holds those buffers until its output is fully consumed.
In a query with several partitions, partitions that finish their input early
start this final merge while other partitions are still buffering input, and
while the operator above the sort is still running. Each early merge takes
whatever the pool admits, and the pool can be exhausted. The partition that is
still sorting then fails to grow its reservation and the query fails with
`ResourcesExhausted`, even though only memory for read-ahead buffers was
taken.
A standalone reproduction (below): 6.67M rows,
`SUM(p) OVER (PARTITION BY g ORDER BY k ROWS UNBOUNDED PRECEDING)` followed by
`MAX`, 10 target partitions, a `FairSpillPool` of 62.5 MB, 40 runs per
setting:
| `sort_spill_reservation_bytes` | main (2c568b47) | this PR |
|---|---|---|
| 1 MB | 25 / 40 succeed | 40 / 40 |
| 1.6 MB | 34 / 40 | 40 / 40 |
| 3 MB | 7 / 40 | 35 / 40 |
The failures are the partitions that are still sorting:
`Failed to allocate additional 3.0 MB for ExternalSorterMerge[6] with 0.0 B
already allocated`, `... 598.8 KB for ExternalSorter[8] with 0.0 B ...`, or
`... 144.1 KB for ExternalSorterMerge[7]` in `BatchBuilder::push_batch` while
merging in-memory batches to spill. With a pool instrumented to record each
consumer's reservation, five partitions in their final merge held 12-29 MB
each when the refusal happened. Capping `max_spill_merge_fan_in` at 4 made
the same runs succeed, which confirms that the merge fan-in causes the failures.
## What changes are included in this PR?
#25383 lets aggregate spill replay leave memory for the aggregate: the merge
asks the pool to admit its buffers plus as much again, then returns the extra.
The sort uses a narrower form of that policy:
- The minimum merge of two runs is admitted as before, against the pool
alone. A sort under tight memory keeps its current behavior: it does not
re-spill runs smaller or emit shorter batches only to leave headroom.
- Each run beyond the minimum is seated only if the pool can admit what the
merge takes beyond the workspace it already holds
(`sort_spill_reservation_bytes`, retained since #24740), plus as much
again. The extra is returned before the merge starts. Seating runs within
the merge's own workspace needs no headroom, because it takes nothing
from others.
When the pool refuses, the merge seats fewer runs and merges in more passes,
as it already does under memory pressure.
The policy is a new `with_fan_in_headroom` option on `StreamingMergeBuilder`
and `MultiLevelMergeBuilder`, enabled by `ExternalSorter`. The aggregate's
replay headroom is unchanged.
### What this does not fix
Each merge leaves as much free as it takes, but merges that start one after
another still compound: in the remaining 3 MB failures, nine partitions held
3-18 MB each (17.9, 11.5, 7.7, 5.1, ... MB), and the last partition to start
could not reserve its own 3 MB `sort_spill_reservation_bytes`. With 10
partitions, the workspaces alone are half of this pool. A guarantee would
need each merge to know how many partitions share the pool, for example a
per-partition share of the limit. I left that out of this PR, because it is
a larger policy change and I would like to hear whether it is wanted.
## What is the testing strategy for this PR?
- `test_final_spilled_merge_leaves_headroom` (sort/spill_tests.rs) spills a
sort and leaves the pool room for the final merge to seat every run, but
not to seat every run and still leave as much free. It checks that whatever
the merge takes beyond its workspace, it leaves as much free, and that the
output is correct and every reservation is released. It fails on main
("the final merge takes 168128 bytes beyond its workspace and leaves 109728
free").
- Existing tests pass unchanged: `datafusion-physical-plan` (2530),
`datafusion-execution`, the `memory_limit` tests in `core_integration`, and
the sort, merge, spilling and aggregate fuzz tests (`extended_tests`).
In particular, the batch-size test under a tight limit
(`..._when_having_to_spill`) and the workspace retention tests from #24740
are unaffected.
<details>
<summary>Reproduction</summary>
`Cargo.toml`:
```toml
[package]
name = "sortrepro"
version = "0.1.0"
edition = "2024"
[dependencies]
datafusion = { git = "https://github.com/apache/datafusion", default-features = false }
tokio = { version = "1", features = ["rt-multi-thread"] }
futures = "0.3"
```
`src/main.rs`, run as `sortrepro <sort_spill_reservation_bytes> <runs>`:
```rust
//! A window over a spilling sort, 10 partitions, one FairSpillPool.
//! Usage: sortrepro <sort_spill_reservation_bytes> <runs>
use std::sync::Arc;
use datafusion::arrow::array::{Decimal256Array, Int32Array, RecordBatch};
use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, i256};
use datafusion::catalog::streaming::StreamingTable;
use datafusion::common::ScalarValue;
use datafusion::functions_aggregate::expr_fn::max;
use datafusion::functions_aggregate::sum::sum_udaf;
use datafusion::logical_expr::expr::WindowFunction;
use datafusion::logical_expr::{
Expr, ExprFunctionExt, WindowFrame, WindowFrameBound, WindowFrameUnits,
WindowFunctionDefinition, col,
};
use datafusion::execution::memory_pool::FairSpillPool;
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::streaming::PartitionStream;
use datafusion::prelude::{SessionConfig, SessionContext};
const ROWS: i32 = 6_666_815;
const BATCH: i32 = 8192;
const POOL: usize = 65_484_800; // 62.5 MB
/// (g, k, p): 1000 groups of increasing k, with a decimal to sum.
#[derive(Debug)]
struct Gen(SchemaRef);
impl PartitionStream for Gen {
fn schema(&self) -> &SchemaRef {
&self.0
}
fn execute(&self, _: Arc<TaskContext>) -> SendableRecordBatchStream {
let schema = Arc::clone(&self.0);
let s = Arc::clone(&schema);
let batches = (0..(ROWS + BATCH - 1) / BATCH).map(move |b| {
let (lo, hi) = (b * BATCH, ((b + 1) * BATCH).min(ROWS));
let k: Vec<i32> = (lo..hi).map(|i| i * 3).collect();
let g: Vec<i32> = k.iter().map(|k| k % 1000).collect();
let p: Vec<i256> = k
.iter()
.map(|k| i256::from_i128((*k as i128 * 7919) % 10_000_000))
.collect();
let p = Decimal256Array::from(p).with_precision_and_scale(76, 2)?;
Ok(RecordBatch::try_new(
Arc::clone(&s),
vec![
Arc::new(Int32Array::from(g)),
Arc::new(Int32Array::from(k)),
Arc::new(p),
],
)?)
});
Box::pin(RecordBatchStreamAdapter::new(
schema,
futures::stream::iter(batches),
))
}
}
async fn run(sort_spill_reservation_bytes: usize) -> datafusion::error::Result<()> {
let mut config = SessionConfig::new()
.with_target_partitions(10)
.with_batch_size(BATCH as usize);
config.options_mut().execution.sort_spill_reservation_bytes =
sort_spill_reservation_bytes;
let runtime = RuntimeEnvBuilder::new()
.with_memory_pool(Arc::new(FairSpillPool::new(POOL)))
.build_arc()?;
let ctx = SessionContext::new_with_config_rt(config, runtime);
let schema = Arc::new(Schema::new(vec![
Field::new("g", DataType::Int32, false),
Field::new("k", DataType::Int32, false),
Field::new("p", DataType::Decimal256(76, 2), false),
]));
let table = StreamingTable::try_new(
Arc::clone(&schema),
vec![Arc::new(Gen(schema))],
)?;
// SELECT max(s) FROM (SELECT sum(p) OVER (PARTITION BY g ORDER BY k
// ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS s FROM t)
let s = Expr::from(WindowFunction::new(
WindowFunctionDefinition::AggregateUDF(sum_udaf()),
vec![col("p")],
))
.partition_by(vec![col("g")])
.order_by(vec![col("k").sort(true, false)])
.window_frame(WindowFrame::new_bounds(
WindowFrameUnits::Rows,
WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
WindowFrameBound::CurrentRow,
))
.build()?
.alias("s");
ctx.read_table(Arc::new(table))?
.window(vec![s])?
.aggregate(vec![], vec![max(col("s"))])?
.collect()
.await?;
Ok(())
}
fn main() {
let args: Vec<String> = std::env::args().collect();
let reservation: usize = args[1].parse().unwrap();
let runs: usize = args[2].parse().unwrap();
let rt = tokio::runtime::Builder::new_multi_thread()
.worker_threads(10)
.enable_all()
.build()
.unwrap();
let mut ok = 0;
for _ in 0..runs {
match rt.block_on(run(reservation)) {
Ok(()) => ok += 1,
Err(e) => println!("{e}"),
}
}
println!("sort_spill_reservation_bytes={reservation}: {ok} / {runs} succeeded");
}
```
</details>
## Are there any user-facing changes?
No API changes. A sort that spills to many runs may merge in more passes when
the pool is shared, instead of failing other consumers of the pool.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#issuecomment-6155723085
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/26264?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `96.15385%` with `1 line` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 82.83%. Comparing base ([`2c568b4`](https://app.codecov.io/gh/apache/datafusion/commit/2c568b478b3aff18de52609c7237f1c3f494d5b1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`a78d433`](https://app.codecov.io/gh/apache/datafusion/commit/a78d43395c94b7392e13f1a5bdb185c2fd722284?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/26264?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...usion/physical-plan/src/sorts/multi\_level\_merge.rs](https://app.codecov.io/gh/apache/datafusion/pull/26264?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Fsorts%2Fmulti_level_merge.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9zb3J0cy9tdWx0aV9sZXZlbF9tZXJnZS5ycw==) | 94.73% | [1 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/26264?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #26264 +/- ##
=======================================
Coverage 82.83% 82.83%
=======================================
Files 1149 1149
Lines 452798 452822 +24
Branches 452798 452822 +24
=======================================
+ Hits 375082 375103 +21
+ Misses 54893 54892 -1
- Partials 22823 22827 +4
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/26264?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "MisterRaindrop (via GitHub) via github" <gi...@datafusion.apache.org>.
MisterRaindrop commented on PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#issuecomment-6675441272
Thanks both. Pushed c7c5b196b.
@jayzhan211 I reproduced your numbers. Headroom is now kept only when the
sort has sibling partitions or another consumer holds pool memory. Single
partition now spills the same rows as without headroom in all four of your
configs. With 4 partitions / 32 MB pool / 1 MB reservation, no headroom
fails 5/5 ("Not enough memory to continue external sort"); with it, 5/5
pass.
@kosiew Added both boundary tests, and applied the wording suggestion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "MisterRaindrop (via GitHub) via github" <gi...@datafusion.apache.org>.
MisterRaindrop commented on PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#issuecomment-6688180513
Thanks @jayzhan211. Pushed 24297ff6c.
Applied all three suggestions and updated the description.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "kosiew (via GitHub) via github" <gi...@datafusion.apache.org>.
kosiew commented on code in PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#discussion_r4452808768
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -1476,6 +1507,71 @@ mod tests {
Ok(())
}
Review Comment:
Could we optionally add a headroom-enabled case with a `5 * batch_memory` pool and no `MergeMemoryPool`? Each run needs `2 * batch_memory`, so please verify that the minimum two runs are admitted, the third is rejected, and `reservation.size()` remains `4 * batch_memory`; the existing tests do not distinguish this minimum-merge exception under a tight budget.
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -611,7 +629,20 @@ impl MultiLevelMergeBuilder {
&& number_of_spills_to_read_for_current_phase
< minimum_number_of_required_streams;
let check_headroom = self.reserve_replay_headroom && !skip_headroom;
- let admission = if check_headroom {
+ let check_fan_in_headroom = !check_headroom
+ && self.reserve_fan_in_headroom
+ && number_of_spills_to_read_for_current_phase
+ >= minimum_number_of_required_streams;
+ let admission = if check_fan_in_headroom {
+ // Ask the pool for what this merge takes beyond its own
+ // workspace, plus as much again, then return the spare bytes.
+ match total_needed.checked_add(total_needed.saturating_sub(owned)) {
Review Comment:
Could we optionally test admission with retained `MergeMemoryPool` workspace `W`, no other live children or loans, and an empty child reservation? For an optional run, cover `T <= W` and `T > W` with parent free space just below and exactly at `2 * (T - W)`, verifying fan-in, child reservation `T` for admitted runs, and final parent charge `max(W, T)` after temporary headroom is released.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "jayzhan211 (via GitHub) via github" <gi...@datafusion.apache.org>.
jayzhan211 commented on code in PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#discussion_r4455032218
##########
datafusion/physical-plan/src/sorts/streaming_merge.rs:
##########
@@ -171,6 +174,13 @@ impl<'a> StreamingMergeBuilder<'a> {
self
}
+ /// Seat spill runs beyond the minimum two only while the pool can admit
+ /// their buffers plus equal headroom. Release the headroom before merging.
Review Comment:
Optional cleanup: the headroom is only as much again as the merge takes beyond its retained workspace, not the full buffers. The same wording is on the field doc in `multi_level_merge.rs` around line 163.
```suggestion
/// Seat spill runs beyond the minimum two only while the pool can admit
/// their buffers plus as much again as they exceed the retained workspace.
/// Release the headroom before merging.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "jayzhan211 (via GitHub) via github" <gi...@datafusion.apache.org>.
jayzhan211 commented on code in PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#discussion_r4500714782
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -372,10 +376,25 @@ impl ExternalSorter {
batch_size,
sort_spill_reservation_bytes,
sort_in_place_threshold_bytes,
+ partitions_reading_input: Arc::new(AtomicUsize::new(1)),
sort_concurrency: Arc::new(OnceLock::new()),
})
}
+ /// Share the count of partitions still reading their input with the
+ /// other partitions of this sort. It must include this partition.
+ fn with_partitions_reading_input(mut self, reading: Arc<AtomicUsize>) -> Self {
+ self.partitions_reading_input = reading;
+ self
+ }
+
+ /// Records that this partition has read all its input, and returns whether
+ /// another partition still reads its own. Executing partitions again
+ /// without `reset_state` can wrap the count, which only keeps headroom.
Review Comment:
Optional cleanup: on the first re-run the count is 0, so this partition gets `false` and keeps no headroom, then the count wraps.
```suggestion
/// Records that this partition has read all its input, and returns whether
/// another partition still reads its own. Executing partitions again
/// without `reset_state` leaves the count wrong, which changes how much
/// headroom a merge keeps but never its output.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "jayzhan211 (via GitHub) via github" <gi...@datafusion.apache.org>.
jayzhan211 commented on code in PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#discussion_r4492889820
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -1463,6 +1525,187 @@ mod tests {
Ok(())
}
+ /// Registers a consumer holding `size` bytes of `pool`, so a merge sharing
+ /// it sees another user and leaves headroom.
+ fn hold(pool: &Arc<dyn MemoryPool>, size: usize) -> Result<MemoryReservation> {
+ let other = MemoryConsumer::new("other consumer").register(pool);
+ other.try_grow(size)?;
+ Ok(other)
+ }
+
+ /// Headroom for others sharing `pool`, with or without sibling partitions.
+ fn headroom(pool: &Arc<dyn MemoryPool>, siblings: bool) -> Option<FanInHeadroom> {
+ Some(FanInHeadroom {
+ pool: Arc::clone(pool),
+ siblings,
+ })
+ }
Review Comment:
Optional cleanup: this helper can return the struct directly, which removes the `.unwrap()` in the first test. The other three callers wrap it in `Some(...)`.
```suggestion
fn headroom(pool: &Arc<dyn MemoryPool>, siblings: bool) -> FanInHeadroom {
FanInHeadroom {
pool: Arc::clone(pool),
siblings,
}
}
```
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -309,6 +309,9 @@ struct ExternalSorter {
/// How much memory to reserve for performing in-memory sort/merges
/// prior to spilling.
sort_spill_reservation_bytes: usize,
+ /// Other partitions of this sort run alongside this one and share its
+ /// memory pool.
+ has_sibling_partitions: bool,
Review Comment:
This stays true after every sibling has finished, so a skewed partition that finishes last still keeps headroom for nobody and rewrites rows. On a window where one of 4 partitions gets all the rows (2M rows, greedy pool, 1 MB reservation, 15 runs), an 80 MB pool goes from 2.00M to 3.30M spilled rows and 88 MB from 2.00M to 3.56M. Pools of 96 MB and up cost nothing.
Counting the partitions still reading their input fixes both of those and keeps most of the gain (49 of 60 runs succeed on the balanced configs in my review comment, against 54 with this PR and 32 without headroom, a difference that could be noise). Merges already in progress are still caught by `pool.reserved() > owned`. The cost is one shared atomic. I ran this with the sort tests and clippy; could you apply it here?
```diff
// ExternalSorter
- has_sibling_partitions: bool,
+ /// Partitions of the owning `SortExec`, this one included, still reading
+ /// their input. Each one will need memory from the pool this sort shares.
+ partitions_reading_input: Arc<AtomicUsize>,
- has_sibling_partitions: false,
+ partitions_reading_input: Arc::new(AtomicUsize::new(1)),
- fn with_sibling_partitions(mut self, has_sibling_partitions: bool) -> Self {
- self.has_sibling_partitions = has_sibling_partitions;
+ fn with_partitions_reading_input(mut self, reading: Arc<AtomicUsize>) -> Self {
+ self.partitions_reading_input = reading;
// ExternalSorter::sort
+ let other_partitions_reading =
+ self.partitions_reading_input.fetch_sub(1, Ordering::AcqRel) > 1;
- self.has_sibling_partitions,
+ other_partitions_reading,
// SortExec: also add the field to `new` and `cloned`, ignore it in the proto
// encoder, and give it a fresh `OnceLock` in `reset_state`
+ /// Partitions still reading their input, shared with their sorters. Set
+ /// on the first `execute`, once the input can no longer change.
+ partitions_reading_input: Arc<OnceLock<Arc<AtomicUsize>>>,
// SortExec::execute
- .with_sibling_partitions(
- self.input.output_partitioning().partition_count() > 1,
- );
+ .with_partitions_reading_input(Arc::clone(
+ self.partitions_reading_input.get_or_init(|| {
+ Arc::new(AtomicUsize::new(
+ self.input.output_partitioning().partition_count(),
+ ))
+ }),
+ ));
```
With this, the existing test sets the counter with `Arc::new(AtomicUsize::new(1 + usize::from(siblings)))`, and the `SortExec` test from my other comment gets a third case: run partition 1 to completion first, then expect partition 0 to write each row once.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1841,7 +1864,10 @@ impl ExecutionPlan for SortExec {
context.session_config().spill_compression(),
&self.metrics_set,
context.runtime_env(),
- )?;
+ )?
+ .with_sibling_partitions(
+ self.input.output_partitioning().partition_count() > 1,
Review Comment:
Nothing tests this line yet: replacing it with `true` or `false` still passes every sort test. Could you add this to `spill_tests.rs` in this PR? I ran it, and it fails under both of those changes. It needs `SortExec`, `ExecutionPlan`, `TestMemoryExec`, `TaskContext` and `SessionConfig` imported.
```rs
/// Whether a `SortExec` keeps fan-in headroom follows its input partitions. The
/// pool fits every run's buffers but not the headroom on top, so with headroom
/// the final merge seats fewer runs and rewrites rows in an extra pass. Only
/// partition 0 runs here, so nothing else ever holds pool memory and the
/// outcome is deterministic.
#[rstest::rstest]
#[case::lone_partition(1, false)]
#[case::idle_sibling(2, true)]
#[tokio::test]
async fn test_sort_exec_keeps_headroom_only_with_sibling_partitions(
#[case] input_partitions: usize,
#[case] rewrites_rows: bool,
) -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("key", DataType::Int64, false),
Field::new("payload", DataType::Utf8, false),
]));
let (batches, rows) = (64, 1024);
let batches = (0..batches)
.map(|batch| {
let keys = (0..rows).map(|row| ((batch * rows + row) * 2654435761) % 1000003);
let payload = (0..rows).map(|row| format!("{row:0120}"));
let columns: Vec<ArrayRef> = vec![
Arc::new(Int64Array::from_iter_values(keys)),
Arc::new(StringArray::from_iter_values(payload)),
];
Ok(RecordBatch::try_new(Arc::clone(&schema), columns)?)
})
.collect::<Result<Vec<_>>>()?;
let total_rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
let mut partitions = vec![batches];
partitions.resize(input_partitions, vec![]);
let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?;
let ordering: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new(
"key", 0,
)))]
.into();
let sort = SortExec::new(ordering, input).with_preserve_partitioning(true);
let runtime = RuntimeEnvBuilder::new()
.with_memory_pool(Arc::new(GreedyMemoryPool::new(3500 * 1024)))
.build_arc()?;
let context = Arc::new(
TaskContext::default()
.with_session_config(
SessionConfig::new()
.with_batch_size(1024)
.with_sort_spill_reservation_bytes(128 * 1024),
)
.with_runtime(runtime),
);
let output: Vec<RecordBatch> = sort.execute(0, context)?.try_collect().await?;
assert_eq!(
output.iter().map(RecordBatch::num_rows).sum::<usize>(),
total_rows
);
let spilled_rows = sort.metrics().unwrap().spilled_rows().unwrap();
if rewrites_rows {
assert!(spilled_rows > total_rows, "{spilled_rows} <= {total_rows}");
} else {
assert_eq!(spilled_rows, total_rows);
}
Ok(())
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] fix: leave memory for others when a spilled sort merges its runs [datafusion]
Posted by "MisterRaindrop (via GitHub) via github" <gi...@datafusion.apache.org>.
MisterRaindrop commented on PR #26264:
URL: https://github.com/apache/datafusion/pull/26264#issuecomment-6798909235
> Ran this locally: slt, the extended sort/spill/window/aggregate fuzz tests and memory_limit all pass, and an A/B against the merge base showed no timing or spill difference for in-memory and spilling sorts.
>
> One question: `partitions_reading_input` only goes down in `sort()`, so a partition whose stream is dropped early stays counted. With three partitions, dropping 0 and finishing 1 and 2 leaves it at 1, and later merges keep headroom for nobody. Output is fine. Worth a small drop guard, here or as a follow up?
>
> Nit: `fetch_sub` can wrap on re-execute without `reset_state`; a saturating `fetch_update` would avoid that.
do it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org