comphead commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4329229999
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -728,6 +784,86 @@ 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_size) = self.sort_run_sizes(&runs)?;
+ let shared_input =
+ shared_input_size.map(|size|
Arc::new(self.reservation.split(size)));
+ 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 {
Review Comment:
`with_lookahead` looks equivalent to
`stream.map(std::future::ready).buffered(2)` wrapped in
`RecordBatchStreamAdapter`. `Buffered` fills its queue to 2 before yielding, so
each yielded batch has exactly one successor already pulled, same as the
`unfold` state here. That would drop the hand-written state machine. I have not
run it, so please confirm the reservation release order in
`test_sort_merge_lookahead_returns_cursor_workspace`.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -2345,6 +2536,119 @@ mod tests {
use insta::assert_snapshot;
use tokio::sync::Notify;
+ #[tokio::test]
+ async fn sort_run_records_compute_before_consumption() -> Result<()> {
+ let batch = RecordBatch::try_from_iter([(
+ "a",
+ Arc::new(Int32Array::from_iter_values((0..4096).rev())) as
ArrayRef,
+ )])?;
+ let sorter = ExternalSorter::new(
+ 0,
+ batch.schema(),
+ [PhysicalSortExpr::new_default(col("a", &batch.schema())?)].into(),
+ 1024,
+ 0,
+ 0,
+ SpillCompression::Uncompressed,
+ &ExecutionPlanMetricsSet::new(),
+ Arc::new(RuntimeEnv::default()),
+ )?;
+ let elapsed_compute = sorter.metrics.baseline.elapsed_compute();
+ for spawned in [false, true] {
+ let reservation = sorter.reservation.new_empty();
+
reservation.try_grow(get_reserved_bytes_for_record_batch(&batch)?)?;
+ let before = elapsed_compute.value();
+ let sort = sorter.sort_run(batch.clone(), reservation, None)?;
+ assert_eq!(elapsed_compute.value(), before);
+
+ let stream = if spawned {
+ ExternalSorter::prepare_sort_run(sort,
Arc::new(Semaphore::new(1)))
+ .await?
+ } else {
+ sort.await?
+ };
+ assert!(elapsed_compute.value() > before);
+ drop(stream);
+ }
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn sort_preparation_holds_permit_and_releases_on_drop() {
+ let permits = Arc::new(Semaphore::new(1));
+ let held = permits.acquire().await.unwrap();
+ let started = Arc::new(Notify::new());
+ let sort = {
+ let started = Arc::clone(&started);
+ async move {
+ started.notify_one();
+
futures::future::pending::<Result<SendableRecordBatchStream>>().await
+ }
+ };
+ let mut preparation =
+ ExternalSorter::prepare_sort_run(sort,
Arc::clone(&permits)).boxed();
+ assert_is_pending(&mut preparation);
+ assert_eq!(Handle::current().metrics().num_alive_tasks(), 0);
+ drop(held);
+ assert_is_pending(&mut preparation);
+ started.notified().await;
+ assert_eq!(permits.available_permits(), 0);
+ drop(preparation);
+ let _permit =
+ tokio::time::timeout(std::time::Duration::from_secs(10),
permits.acquire())
+ .await
+ .expect("cancelled preparation must release its permit")
+ .unwrap();
+ }
+
+ #[tokio::test]
+ async fn sort_preparation_propagates_errors_and_releases_permit() {
+ let permits = Arc::new(Semaphore::new(1));
+ let sort = async { Err(internal_datafusion_err!("sort failed")) };
+ let error = ExternalSorter::prepare_sort_run(sort,
Arc::clone(&permits))
+ .await
+ .err()
+ .expect("preparation must propagate the sort error");
+ assert!(error.to_string().contains("sort failed"), "{error}");
+ assert_eq!(permits.available_permits(), 1);
+ }
+
+ #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+ async fn sort_execution_initializes_and_resets_concurrency() -> Result<()>
{
Review Comment:
This test collects but never checks the output. Since the claim is that the
spawned path matches the inline one, please assert the collected rows against
the same plan run on a current-thread runtime. Use a payload column so equal
keys are distinguishable.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -728,6 +784,86 @@ 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_size) = self.sort_run_sizes(&runs)?;
Review Comment:
Two small points. (1) These three lines (`sort_run_sizes`, shared-input
split, per-run `split`) are repeated in `sort_run_streams`. A shared helper
returning the per-run reservations plus the optional shared one would keep the
two paths from drifting. (2) `prepare_sort_runs` takes `&mut self` but only
calls `reservation.split(&self)`, like `sort_run_streams(&self)`. `&self` is
enough if the future stays `Send` (not checked).
##########
datafusion/physical-plan/src/sorts/sort/spill_tests.rs:
##########
@@ -533,6 +543,155 @@ async fn
test_chunked_string_view_spill_can_use_workspace() -> Result<()> {
check_chunked_string_view_workspace(true).await
}
+/// Prepared runs retain their workspace loans until merging starts. Looking
+/// ahead during the merge returns enough workspace for its batch and cursors.
+async fn check_sort_preparation_workspace(consume_output: bool) -> Result<()> {
Review Comment:
The description says the merge receives runs in original order so equal-key
ties are preserved, but nothing pins that. Both runs here are the same batch
and only have the key column, so any run order gives identical output. A test
with a payload column, where the earlier run finishes last (for example a much
larger first run), would cover the `try_buffer_unordered` plus `streams[index]`
reassembly.
--
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]
