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: [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]
