geoffreyclaude commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4279842931
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -716,6 +752,73 @@ impl ExternalSorter {
.build()
}
+ /// Prepare runs concurrently and return their streams in the original run
order.
+ 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 jobs = runs.into_iter().enumerate().map(move |(index, batch)| {
+ let reservation = self
+ .reservation
+ .split(get_reserved_bytes_for_record_batch(&batch)?);
+ let stream = self.sort_batch_stream(batch, reservation)?;
+ let permits = Arc::clone(&concurrency.permits);
+ Ok::<_, DataFusionError>(async move {
+ let stream = Self::prepare_sort_run(stream, permits).await?;
+ Ok((index, stream))
+ })
+ });
+ let mut jobs =
+
futures::stream::iter(jobs).try_buffer_unordered(concurrency.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 through its first batch and release its slot before
merging.
+ async fn prepare_sort_run(
+ mut stream: SendableRecordBatchStream,
+ 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}"))?;
+ SpawnedTask::spawn(async move {
+ // Match the buffered producer's first-batch handoff so spill
+ // workspace can be reused while the remaining runs are sorted.
+ let first = stream.try_next().await?.map(Ok);
+ let schema = stream.schema();
+ let ready = Box::pin(RecordBatchStreamAdapter::new(
+ schema,
+ futures::stream::unfold(
+ (first, stream),
+ |(buffered, mut stream)| async move {
+ let batch = buffered?;
+ // Sorting materializes every chunk on the first poll.
+ // Refill the one-batch buffer before handing off its
+ // current batch, returning workspace before the merge
+ // reserves memory for that batch and its cursor.
+ let next = stream.next().await;
+ Some((batch, (next, stream)))
Review Comment:
Good catch. Preparation now keeps all sorted output reserved until merging
starts. During merging, the one-batch lookahead still temporarily releases the
prefetched batch’s reservation so the merge can reuse that workspace.
--
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]
