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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
