zhuqi-lucas opened a new issue, #27003:
URL: https://github.com/apache/datafusion/issues/27003
### Is your feature request related to a problem or challenge?
`SortPreservingMergeExec` only polls the input that currently holds the
smallest key, and it does not spawn tasks or read ahead. When its inputs are
not buffered and their key ranges are disjoint (for example, files written with
range-based splitting, so file 0 holds keys `A..C`, file 1 holds `C..F`, and so
on), the merge drains one input completely before it pulls from the next one.
Every input below the merge, including its scan, runs one after another.
A common plan where this happens is a two-phase aggregate over files that
are each sorted by the group key:
```
AggregateExec: mode=Final
SortPreservingMergeExec: [k ASC]
AggregateExec: mode=Partial (ordering_mode=Sorted)
DataSourceExec: one file group per file
```
`add_merge_on_top` picks `SortPreservingMergeExec` because the Partial
output is ordered. With N files on an object store, the query takes about N
times the time to read one file.
The only place `BufferExec` is put under a merge today is `PushdownSort`,
when it removes a `SortExec` below `SortPreservingMergeExec`. That does not
cover this plan, which never had a `SortExec`, and it is skipped when
`enable_sort_pushdown = false`.
### Repro
On DataFusion 55.2.0: 32 parquet files, 1M rows each, each sorted by `(k,
ts, seq)`. A `ThrottledStore` adds 45 ms to every GET. The query is:
```sql
SELECT k, last_value(v1 ORDER BY ts, seq), last_value(v2 ORDER BY ts, seq)
FROM t GROUP BY k ORDER BY k
```
with `repartition_aggregations = false`, `prefer_existing_sort = true`, and
`repartition_file_scans = false`.
| file key ranges | plan | 45 ms per GET | no added latency |
|---|---|---|---|
| disjoint | as planned | **26.24 s** | 0.58 s |
| disjoint | `repartition_aggregations = true`, `prefer_existing_sort =
false` | **21.72 s** | 0.50 s |
| disjoint | `BufferExec` wrapped around the merge input | **0.94 s** | 0.11
s |
| overlapping | as planned | 1.15 s | 0.09 s |
| overlapping | `repartition_aggregations = true`, `prefer_existing_sort =
false` | 0.93 s | 0.09 s |
| overlapping | `BufferExec` wrapped around the merge input | 0.88 s | 0.10
s |
All runs return the same 32,000 rows in key order. 26.24 s / 45 ms is about
580 GETs issued one at a time, which is all 32 files read in sequence. With
overlapping ranges the merge rotates across inputs, so this does not show up.
On local disk without latency the gap is small, so the problem mainly shows up
on object stores.
`repartition_aggregations = true` together with `prefer_existing_sort =
false` does not avoid it in this plan. The hash `RepartitionExec` keeps
`preserve_order=true`, because `FinalPartitioned` has no ordering requirement
and maintains its input order (the no-ordering-requirement branch at the end of
`ensure_distribution`). Each output partition then merges all inputs in key
order:
```
SortPreservingMergeExec: [k@0 ASC NULLS LAST]
AggregateExec: mode=FinalPartitioned, gby=[k@0 as k], ordering_mode=Sorted
RepartitionExec: partitioning=Hash([k@0], 32), input_partitions=32,
preserve_order=true
AggregateExec: mode=Partial, gby=[k@0 as k], ordering_mode=Sorted
DataSourceExec: file_groups={32 groups}
```
### Describe the solution you'd like
Put a `BufferExec` under `SortPreservingMergeExec` when its inputs are not
already buffered, as a rule (or setting) that does not depend on sort pushdown.
Each input then runs ahead in its own task and the merge reads from memory.
Conditions to consider:
- skip it when the merge has a `fetch`, because buffering ahead would read
much more than a top-k needs
- bound memory: with N inputs the buffers can total N times the capacity
(#21440 proposes a row limit plus a memory cap)
- `BufferExec` returns an error when `try_grow` fails, instead of waiting,
so a tight memory pool could fail queries that pass today
A narrower first step would be to buffer only when the merge input reduces
data heavily, such as a Partial aggregate.
### Describe alternatives you've considered
- Parallel merge inside `SortPreservingMergeExec` (#21381). That addresses
merge CPU after sort elimination. With disjoint ranges every merge node still
pulls one child at a time, so read-ahead is what fixes this case.
- Using file statistics to detect disjoint ranges and concatenate the inputs
in order. That avoids the comparisons but is still sequential without
read-ahead.
### Additional context
Related: #21381, #21182, #21440, #24270
<details>
<summary>Repro code</summary>
`Cargo.toml`
```toml
[package]
name = "spm-repro"
version = "0.1.0"
edition = "2024"
[dependencies]
datafusion = "=55.2.0"
object_store = "0.13.2"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
url = "2"
```
Needs rustc 1.94 or later (the DataFusion 55 MSRV).
`src/main.rs`
```rust
//! SortPreservingMergeExec over unbuffered inputs whose key ranges are
disjoint.
//!
//! gen <dir> <files> <keys_per_file> <rows_per_key> <disjoint|overlap>
//! run <dir> <files> <get_latency_ms> <spm|repartition|buffer>
use std::sync::Arc;
use std::time::{Duration, Instant};
use datafusion::arrow::array::{Array, Float64Array, Int64Array, RecordBatch,
StringArray};
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::common::tree_node::{Transformed, TreeNode};
use datafusion::datasource::file_format::parquet::ParquetFormat;
use datafusion::datasource::listing::{
ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
};
use datafusion::error::Result;
use datafusion::parquet::arrow::ArrowWriter;
use datafusion::parquet::file::properties::WriterProperties;
use datafusion::physical_plan::buffer::BufferExec;
use
datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
use datafusion::physical_plan::{ExecutionPlan, collect, displayable};
use datafusion::prelude::{SessionConfig, SessionContext, col};
use object_store::local::LocalFileSystem;
use object_store::throttle::{ThrottleConfig, ThrottledStore};
use url::Url;
#[tokio::main(flavor = "multi_thread", worker_threads = 16)]
async fn main() -> Result<()> {
let a: Vec<String> = std::env::args().collect();
let n = |i: usize| a[i].parse::<usize>().unwrap();
match a[1].as_str() {
"gen" => generate(&a[2], n(3), n(4), n(5), a[6] == "overlap"),
"run" => run(&a[2], n(3), n(4) as u64, &a[5]).await,
_ => panic!("usage"),
}
}
fn schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![
Field::new("k", DataType::Utf8, false),
Field::new("ts", DataType::Int64, false),
Field::new("seq", DataType::Int64, false),
Field::new("v1", DataType::Float64, false),
Field::new("v2", DataType::Float64, false),
]))
}
/// Every file is sorted by (k, ts, seq). `disjoint`: file i holds keys
`K{i}_{j}`, so the
/// files' key ranges follow one another. `overlap`: every file spans the
whole key range.
fn generate(dir: &str, files: usize, keys: usize, rows: usize, overlap:
bool) -> Result<()> {
std::fs::create_dir_all(dir)?;
let props =
WriterProperties::builder().set_max_row_group_row_count(Some(65536)).build();
for f in 0..files {
let file =
std::fs::File::create(format!("{dir}/part-{f:04}.parquet"))?;
let mut w = ArrowWriter::try_new(file, schema(),
Some(props.clone()))?;
for j in 0..keys {
let k = if overlap { format!("K{j:05}_{f:04}") } else {
format!("K{f:04}_{j:05}") };
let batch = RecordBatch::try_new(
schema(),
vec![
Arc::new(StringArray::from(vec![k.as_str(); rows])),
Arc::new(Int64Array::from_iter_values(0..rows as i64)),
Arc::new(Int64Array::from_iter_values(0..rows as i64)),
Arc::new(Float64Array::from_iter_values((0..rows).map(|i| (f + j + i) as f64))),
Arc::new(Float64Array::from_iter_values((0..rows).map(|i| (f + j + i + 1) as
f64))),
],
)?;
w.write(&batch)?;
}
w.close()?;
}
Ok(())
}
async fn run(dir: &str, files: usize, latency_ms: u64, variant: &str) ->
Result<()> {
let mut config = SessionConfig::new().with_target_partitions(files);
let opts = config.options_mut();
opts.optimizer.repartition_file_scans = false;
// `repartition`: let the planner hash-repartition the aggregation and
drop the merge.
let repartition = variant == "repartition";
opts.optimizer.repartition_aggregations = repartition;
opts.optimizer.prefer_existing_sort = !repartition;
let ctx = SessionContext::new_with_config(config);
// Simulate object store latency on every GET.
let throttle = ThrottleConfig {
wait_get_per_call: Duration::from_millis(latency_ms),
..Default::default()
};
let store = ThrottledStore::new(LocalFileSystem::new_with_prefix(dir)?,
throttle);
ctx.register_object_store(&Url::parse("throttle://bucket").unwrap(),
Arc::new(store));
let options = ListingOptions::new(Arc::new(ParquetFormat::default()))
.with_file_extension(".parquet")
.with_file_sort_order(vec![vec![
col("k").sort(true, false),
col("ts").sort(true, false),
col("seq").sort(true, false),
]]);
let table =
ListingTableConfig::new(ListingTableUrl::parse("throttle://bucket/")?)
.with_listing_options(options)
.with_schema(schema());
ctx.register_table("t", Arc::new(ListingTable::try_new(table)?))?;
let sql = "SELECT k, last_value(v1 ORDER BY ts, seq) AS v1,
last_value(v2 ORDER BY ts, seq) AS v2 \
FROM t GROUP BY k ORDER BY k";
let mut plan = ctx.sql(sql).await?.create_physical_plan().await?;
if variant == "buffer" {
plan = buffer_spm_inputs(plan)?;
}
println!("{}", displayable(plan.as_ref()).indent(false));
let start = Instant::now();
let batches = collect(plan, ctx.task_ctx()).await?;
let rows: usize = batches.iter().map(|b| b.num_rows()).sum();
let keys: Vec<String> = batches
.iter()
.flat_map(|b| {
let c =
b.column(0).as_any().downcast_ref::<StringArray>().unwrap();
(0..c.len()).map(|i| c.value(i).to_string()).collect::<Vec<_>>()
})
.collect();
let sorted = keys.windows(2).all(|w| w[0] < w[1]);
println!(
"variant={variant} latency={latency_ms}ms elapsed={:.2}s rows={rows}
sorted={sorted}",
start.elapsed().as_secs_f64()
);
Ok(())
}
/// Wraps every input of a SortPreservingMergeExec in a BufferExec.
fn buffer_spm_inputs(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn
ExecutionPlan>> {
plan.transform_up(|node| {
let Some(spm) = node.downcast_ref::<SortPreservingMergeExec>() else {
return Ok(Transformed::no(node));
};
let buffered = Arc::new(BufferExec::new(Arc::clone(spm.input()), 1
<< 30));
let new_spm = SortPreservingMergeExec::new(spm.expr().clone(),
buffered).with_fetch(spm.fetch());
Ok(Transformed::yes(Arc::new(new_spm) as Arc<dyn ExecutionPlan>))
})
.map(|t| t.data)
}
```
```
cargo run --release -- gen data 32 1000 1000 disjoint
cargo run --release -- run data 32 45 spm
cargo run --release -- run data 32 45 repartition
cargo run --release -- run data 32 45 buffer
```
</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: [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]
