geoffreyclaude commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4279837462
##########
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.
Review Comment:
indeed, removed.
--
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]
