jayshrivastava commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4286287268
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -822,13 +955,26 @@ impl ExternalSorter {
Arc::new(self.reservation.split(retained))
});
Review Comment:
It's odd that in the shared case, this reserves memory for all the runs
together here. In the non shared case, the caller calls
`self.reservation.split(retained)`.
Maybe this function should just return an `usize` for `shared_input` so the
caller is always responsible?
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -793,22 +927,21 @@ impl ExternalSorter {
.map(|_| RecordBatchMemoryCounter::new())
}
- /// Sorts each run into its own stream, splitting `self.reservation`
between
- /// them.
+ /// Calculates each run's reservation size and reserves any shared input.
Review Comment:
```suggestion
/// Calculates each run's reservation size and reserves any shared input.
///
/// A run's reservation size is the size of the RecordBatch (excluding
shared buffers, which are
/// accounted in the 2nd return value) + the size of creating a copy of
the RecordBatch, which is
/// space for the sorted output.
```
##########
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:
Oh neat, with_elapsed_compute is quite handy
##########
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:
My earlier comment was actually about the lookahead entirely.
You have potentially `lookahead * total_runs` amount of memory untracked
here where `total_runs` may be more than the number of permits. This exists
because we want to free up space borrowed from the merger pool before the
merger starts - but why not solve this in the merger without using a lookahead?
Ex. have the merger poll the stream (which frees space in its pool) then do the
merge or allocate more space?
This can also be follow up work because your PR is a strict improvement over
the previous code. Just wanted to point this out bc it looks like an accounting
issue.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -793,22 +927,21 @@ impl ExternalSorter {
.map(|_| RecordBatchMemoryCounter::new())
}
- /// Sorts each run into its own stream, splitting `self.reservation`
between
- /// them.
+ /// Calculates each run's reservation size and reserves any shared input.
///
/// Runs can share buffers, for example zero-copy slices of one larger
batch,
/// and `self.reservation` covers each shared buffer once. A shared buffer
/// stays alive until the last run holding it is sorted, so if any buffers
are
/// shared, the reservation for every run's input is held until all the
runs
/// are sorted. Otherwise each run releases its input once it is sorted.
- fn sort_run_streams(
+ fn sort_run_reservations(
Review Comment:
Hypothetically, could there be a record batch outside the runs which is also
shared?
That's fine but might be worth commenting about it. It means we may be
overaccounting.
--
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]
