jayshrivastava commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4274617231
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -690,16 +704,38 @@ impl ExternalSorter {
batches
};
- let streams = runs
- .into_iter()
- .map(|batch| {
- let reservation = self
- .reservation
- .split(get_reserved_bytes_for_record_batch(&batch)?);
- let input = self.sort_batch_stream(batch, reservation)?;
- Ok(spawn_buffered(input, 1))
- })
- .collect::<Result<_>>()?;
+ let sort_concurrency = Arc::clone(&self.sort_concurrency);
+ let concurrency = Handle::try_current()
+ .ok()
+ .filter(|handle| handle.runtime_flavor() ==
RuntimeFlavor::MultiThread)
+ .map(|handle| {
+ sort_concurrency.get_or_init(|| {
+ let limit = handle.metrics().num_workers();
+ SortConcurrency {
+ limit,
+ permits: Arc::new(Semaphore::new(limit)),
+ }
+ })
+ });
+ let streams = if let Some(concurrency) = concurrency {
+ // Exclude scheduling waits from elapsed compute.
+ timer.stop();
+ let streams = self
+ .prepare_sort_runs(runs, concurrency)
+ .with_elapsed_compute(elapsed_compute.clone())
+ .await?;
+ timer.restart();
Review Comment:
So pausing the timer actually fails to capture the time spent in the
`sort_batch_stream` stream which calls `sort_batch_chunked`, which is elapsed
compute we want to capture. Maybe pass the timer as a param? Or something
better, a timer param feels like a smell.
##########
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:
Haven't read the consumer side of this stream yet but here's one observation
-
Since the batch is unaccounted, why not re-reserve the memory here in
`self.merge_pool`?
- In total across all permits, there's `batch_size * semaphore size`, so
that's a lot of unaccounted memory right?
- The consumed stream returns all the reserved memory (to the workspace or
otherwise) before yielding a batch
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -88,6 +91,15 @@ struct ExternalSorterMetrics {
spill_metrics: SpillMetrics,
}
+/// Shares finite batch-sort jobs across the partitions of one [`SortExec`].
+#[derive(Debug)]
+struct SortConcurrency {
+ /// Bounds preparation futures buffered by each partition.
+ limit: usize,
Review Comment:
Hmm I guess you need this bc there's no API for total semaphore size?
##########
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:
I don't think this comment adds much value.
--
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]
