geoffreyclaude commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4334524277
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -728,6 +784,86 @@ 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 mut streams: Vec<_> = (0..runs.len()).map(|_| None).collect();
+ 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 mut jobs = futures::stream::iter(jobs)
+ .try_buffer_unordered(concurrency.per_partition_limit);
+ while let Some((index, stream)) = jobs.try_next().await? {
+ streams[index] = Some(stream);
+ }
+ streams
+ .into_iter()
+ .collect::<Option<Vec<_>>>()
+ .ok_or_else(|| internal_datafusion_err!("Missing prepared sort
run"))
+ }
+
+ /// Sort one run under a shared permit, retaining its output reservations.
+ ///
+ /// The sort must finish independently of other runs and partitions so its
+ /// permit can be returned before the merge starts consuming output.
+ async fn prepare_sort_run(
+ sort: impl Future<Output = Result<SendableRecordBatchStream>> + Send +
'static,
+ permits: Arc<Semaphore>,
+ ) -> Result<SendableRecordBatchStream> {
+ // Acquire before spawning so queued runs do not become Tokio tasks.
+ let permit = permits
+ .acquire_owned()
+ .await
+ .map_err(|e| internal_datafusion_err!("Sort concurrency closed:
{e}"))?;
+ let stream = SpawnedTask::spawn(async move {
+ // Keep the permit until this job exits, even if its caller is
cancelled.
+ let _permit = permit;
+ sort.await
+ })
+ .join_unwind()
+ .await
+ .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
+ Ok(Self::with_lookahead(stream))
+ }
+
+ /// [`Self::sort_run`] materializes all output and releases each batch's
+ /// reservation when that batch is pulled. Reading one batch ahead returns
+ /// extra spill workspace for the merge's current batch and cursor.
+ /// Each run can hold one prefetched batch whose reservation is already
released.
+ fn with_lookahead(stream: SendableRecordBatchStream) ->
SendableRecordBatchStream {
Review Comment:
You're right: this was premature optimization (I was concerned about extra
per-batch allocations and bookkeeping). I've switched to `buffered(2)`: it's
much clearer, and the query-level benchmarks showed no performance difference.
I also ran `test_sort_merge_lookahead_returns_cursor_workspace` and can confirm
it passes with `buffered(2)`.
--
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]
