comphead commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4339030289
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -728,6 +784,87 @@ impl ExternalSorter {
.build()
}
+ /// Sort runs concurrently, preserving their order for the merge's
equal-key ties.
+ async fn prepare_sort_runs(
+ &mut self,
+ runs: Vec<RecordBatch>,
+ concurrency: &SortConcurrency,
+ ) -> Result<Vec<SendableRecordBatchStream>> {
+ let run_count = runs.len();
+ let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?;
+ let shared_input =
+ shared_input_size.map(|size|
Arc::new(self.reservation.split(size)));
+ let jobs =
+ runs.into_iter()
+ .zip(sizes)
+ .enumerate()
+ .map(move |(index, (batch, size))| {
+ let reservation = self.reservation.split(size);
+ let sort = self.sort_run(batch, reservation,
shared_input.clone())?;
+ let permits = Arc::clone(&concurrency.permits);
+ Ok::<_, DataFusionError>(async move {
+ let stream = Self::prepare_sort_run(sort,
permits).await?;
+ Ok((index, stream))
+ })
+ });
+ // Ordered buffering would keep completed runs behind a slow earlier
run,
+ // filling the window and preventing new jobs from starting.
+ let jobs = futures::stream::iter(jobs)
+ .try_buffer_unordered(concurrency.per_partition_limit);
Review Comment:
One job per run is what makes the per-partition window, the unordered
buffering, the `(index, stream)` reassembly and `collect_prepared_runs`
necessary. An alternative is to split the runs into at most
`per_partition_limit` contiguous, byte-balanced chunks, with one job sorting
each chunk sequentially.
- All jobs are in flight at once, so ordered `try_buffered` has no
head-of-line blocking and run order is preserved without index bookkeeping.
- Per-job overhead (permit, spawn, join, boxed stream) is paid per chunk
rather than per run. With 32-row runs that overhead is likely comparable to the
sort itself. I have not measured this.
- The shared semaphore still bounds jobs across partitions.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -707,11 +754,20 @@ impl ExternalSorter {
batches
};
- let streams = self
- .sort_run_streams(runs)?
- .into_iter()
- .map(|stream| spawn_buffered(stream, 1))
- .collect();
+ let concurrency =
SortConcurrency::for_current_runtime(&self.sort_concurrency);
Review Comment:
The bound caps the task count, but I think the root cause is the run count.
187k runs is SF1 `lineitem` at one run per 32-row batch. With this PR the merge
still has 187k streams, each with its own lookahead batch (the gap tracked in
#26411).
The run count stays at one per input batch because
`coalesce_in_mem_batches_into_runs` only runs on the output path and only for
single-column sorts (`coalesce_runs && self.expr.len() == 1` just above).
Multi-column sorts and every spill keep one run per input batch.
Would a minimum run size, for example coalescing runs below `batch_size`
rows regardless of column count, fix the reported case more directly? It would
cut task count, merge fan-in and lookahead memory together. For comparison, the
description reports at most 123 ms for 8192-row batches on `main` and 155 to
165 ms for this PR on 32-row batches. Those are separate runs, so this is only
indicative. The semaphore would still make sense as a backstop for very large
inputs. This could also be a follow-up if you prefer to keep this PR focused.
--
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]
