geoffreyclaude commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4292315507
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -728,6 +784,84 @@ 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) = self.sort_run_reservations(&runs)?;
+ 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 {
+ let schema = stream.schema();
+ let batches = futures::stream::unfold(
+ (None, stream),
+ |(lookahead, mut stream)| async move {
+ let current = match lookahead {
+ Some(batch) => batch,
+ None => stream.next().await?,
+ };
+ let next = stream.next().await;
+ Some((current, (next, stream)))
+ },
+ );
+ Box::pin(RecordBatchStreamAdapter::new(schema, batches))
Review Comment:
You’re right: the gap can extend to one prefetched batch per run. The
previous buffered producers had the same issue, so I opened #26411 to track the
accounting fix separately.
--
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]
