From 986d14c9ddd9f629817254d0419d8ef59809dd92 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Thu, 3 Sep 2026 22:53:44 +0800 Subject: [PATCH 1/4] feat: support async spill writers Object store spill backends previously had to block while uploading data. Add an async writer path so external sort can stream multipart uploads without buffering complete spill files in memory. CLOSES #23247 --- .../examples/data_io/object_store_spill.rs | 242 +++++-- datafusion/execution/src/lib.rs | 2 +- datafusion/execution/src/spill_file.rs | 80 ++- datafusion/physical-plan/src/sorts/sort.rs | 215 +++++- .../src/spill/in_progress_spill_file.rs | 146 ++++- datafusion/physical-plan/src/spill/mod.rs | 615 +++++++++++++++--- .../physical-plan/src/spill/spill_manager.rs | 145 ++++- 7 files changed, 1244 insertions(+), 201 deletions(-) diff --git a/datafusion-examples/examples/data_io/object_store_spill.rs b/datafusion-examples/examples/data_io/object_store_spill.rs index d7d5392f66953..c714c45570a64 100644 --- a/datafusion-examples/examples/data_io/object_store_spill.rs +++ b/datafusion-examples/examples/data_io/object_store_spill.rs @@ -24,26 +24,31 @@ //! See [`datafusion::execution::memory_pool`] for more information on how //! DataFusion decides when operators should spill, and [`SpillFile`] for the //! spill file abstraction this example implements. -use std::future::Future; -use std::io::Write; +//! +//! This example exercises the asynchronous external-sort spill path. Execution +//! paths that require a partially written local file to be readable still use +//! the synchronous spill API. use std::path::Path as StdPath; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use async_trait::async_trait; use bytes::Bytes; -use datafusion::common::Result; +use datafusion::common::{Result, not_impl_err}; use datafusion::execution::disk_manager::DiskManagerBuilder; use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::execution::{SpillFile, SpillWriter, TempFileFactory}; +use datafusion::execution::{AsyncSpillWriter, SpillFile, SpillWriter, TempFileFactory}; use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_common::exec_err; use futures::{Stream, StreamExt, TryStreamExt, stream}; use object_store::local::LocalFileSystem; use object_store::path::Path; -use object_store::{ObjectStore, ObjectStoreExt, PutPayload}; +use object_store::{MultipartUpload, ObjectStore, ObjectStoreExt, PutPayloadMut}; use tempfile::tempdir; +/// Most remote object stores require non-final multipart parts to be at least 5 MiB. +const OBJECT_STORE_PART_SIZE: usize = 5 * 1024 * 1024; + /// Demonstrates configuring DataFusion with spill files backed by an ObjectStore. pub async fn object_store_spill() -> Result<()> { // A real system would use S3, GCS, Azure, or some other ObjectStore for @@ -54,7 +59,7 @@ pub async fn object_store_spill() -> Result<()> { Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path())?); // Create the custom TempFileFactory that creates spill files in the ObjectStore. - let temp_file_factory = Arc::new(ObjectStoreTempFileFactory::new(store)); + let temp_file_factory = Arc::new(ObjectStoreTempFileFactory::new(Arc::clone(&store))); let disk_manager_builder = DiskManagerBuilder::default().with_temp_file_factory(temp_file_factory.clone()); let runtime = RuntimeEnvBuilder::new() @@ -94,6 +99,19 @@ pub async fn object_store_spill() -> Result<()> { temp_file_factory.created_files() > 0, "expected the custom TempFileFactory to be used for spilling" ); + // Ensure the workload crosses a multipart boundary so the example uploads + // at least one complete part before the final `finish` call. + let spill_prefix = Path::from("spill"); + let spill_objects = store + .list(Some(&spill_prefix)) + .try_collect::>() + .await?; + assert!( + spill_objects + .iter() + .any(|object| object.size >= OBJECT_STORE_PART_SIZE as u64), + "expected at least one spill object to exceed the multipart part size" + ); Ok(()) } @@ -154,8 +172,8 @@ impl TempFileFactory for ObjectStoreTempFileFactory { /// Logical spill file stored at an ObjectStore path. /// -/// DataFusion writes spill data by calling [`SpillFile::open_writer`] and reads -/// it back by calling [`SpillFile::read_stream`]. +/// DataFusion writes spill data by calling [`SpillFile::open_async_writer`] and +/// reads it back by calling [`SpillFile::read_stream`]. struct ObjectStoreSpillFile { /// ObjectStore containing the spill object. store: Arc, @@ -165,6 +183,7 @@ struct ObjectStoreSpillFile { size: Arc, } +#[async_trait] impl SpillFile for ObjectStoreSpillFile { /// Return no local filesystem path because the spill file is accessed through ObjectStore. fn path(&self) -> Option<&StdPath> { @@ -193,81 +212,180 @@ impl SpillFile for ObjectStoreSpillFile { Ok(Box::pin(stream)) } - /// Open a synchronous writer for this spill file. + /// This example backend supports the asynchronous external-sort spill path only. fn open_writer(&self) -> Result> { - // Create a writer that buffers bytes and uploads them on finish. + not_impl_err!("Synchronous spill writing is not supported by this backend") + } + + /// Open an asynchronous, multipart-capable writer for this spill file. + async fn open_async_writer(&self) -> Result> { + let upload = self.store.put_multipart(&self.location).await?; Ok(Box::new(ObjectStoreSpillWriter { - store: Arc::clone(&self.store), - location: self.location.clone(), + upload, + buffer: PutPayloadMut::new(), size: Arc::clone(&self.size), - buffer: Vec::new(), + bytes_written: 0, })) } } -/// Adapts DataFusion's [`SpillWriter`] API to ObjectStore. -/// -/// This simple example buffers bytes in memory and uploads them in -/// [`SpillWriter::finish`]. A production remote implementation should consider -/// multipart or streaming uploads. +/// Adapts DataFusion's [`AsyncSpillWriter`] API to ObjectStore. struct ObjectStoreSpillWriter { - /// ObjectStore to read/write bytes to. - store: Arc, - /// ObjectStore path to upload to. - location: Path, + /// Multipart upload used to stream completed parts to the store. + upload: Box, + /// Buffers at most one part while preserving owned `Bytes` chunks. + buffer: PutPayloadMut, /// Shared size field on the corresponding [`ObjectStoreSpillFile`]. size: Arc, - /// Buffered spill bytes waiting to be uploaded. - /// - /// This simple example buffers the spill and uploads it on finish. - /// Production remote stores should consider multipart or streaming uploads. - buffer: Vec, + /// Number of bytes passed to the writer. + bytes_written: u64, +} + +impl ObjectStoreSpillWriter { + async fn flush_part(&mut self) -> object_store::Result<()> { + if self.buffer.is_empty() { + return Ok(()); + } + + let part = std::mem::take(&mut self.buffer).freeze(); + self.upload.put_part(part).await + } } -impl Write for ObjectStoreSpillWriter { - /// Append bytes to the in-memory buffer. - fn write(&mut self, buf: &[u8]) -> std::io::Result { - // Buffer bytes written through the synchronous Write API. - self.buffer.extend_from_slice(buf); - Ok(buf.len()) +#[async_trait] +impl AsyncSpillWriter for ObjectStoreSpillWriter { + async fn write_all(&mut self, mut data: Bytes) -> Result<()> { + let len = data.len() as u64; + while !data.is_empty() { + let remaining = OBJECT_STORE_PART_SIZE - self.buffer.content_length(); + if data.len() < remaining { + // A Bytes slice can pin an entire Arrow allocation. Copy only + // the tail retained after this call so it remains accurately + // bounded after DataFusion releases the batch reservation. + self.buffer.push(Bytes::copy_from_slice(&data)); + break; + } + + self.buffer.push(data.split_to(remaining)); + self.flush_part().await?; + } + self.bytes_written += len; + Ok(()) + } + + async fn finish(&mut self) -> Result<()> { + self.flush_part().await?; + self.upload.complete().await?; + self.size.store(self.bytes_written, Ordering::Relaxed); + Ok(()) } - /// No-op because data is committed in [`SpillWriter::finish`]. - fn flush(&mut self) -> std::io::Result<()> { + async fn abort(&mut self) -> Result<()> { + // Release locally buffered spill data before waiting on remote cleanup. + self.buffer = PutPayloadMut::new(); + self.upload.abort().await?; Ok(()) } } -impl SpillWriter for ObjectStoreSpillWriter { - /// Upload buffered bytes to ObjectStore and mark the spill file complete. - fn finish(&mut self) -> Result<()> { - // Move the buffered bytes into the upload future. - let store = Arc::clone(&self.store); - let location = self.location.clone(); - let data = std::mem::take(&mut self.buffer); - let size = data.len() as u64; - - // This simple example buffers the spill and uploads it on finish. - // Production remote stores should consider multipart or streaming uploads. - block_on_object_store(async move { - store - .put(&location, PutPayload::from_bytes(data.into())) - .await?; +#[cfg(test)] +mod tests { + use super::*; + use object_store::{PutPayload, PutResult, UploadPart}; + use std::sync::Mutex; + use std::sync::atomic::AtomicBool; + + #[derive(Debug)] + struct RecordingMultipartUpload { + part_sizes: Arc>>, + } + + #[async_trait] + impl MultipartUpload for RecordingMultipartUpload { + fn put_part(&mut self, data: PutPayload) -> UploadPart { + self.part_sizes.lock().unwrap().push(data.content_length()); + Box::pin(async { Ok(()) }) + } + + async fn complete(&mut self) -> object_store::Result { + Ok(PutResult { + e_tag: None, + version: None, + }) + } + + async fn abort(&mut self) -> object_store::Result<()> { Ok(()) - })?; + } + } + + struct DropTrackingOwner { + data: Vec, + dropped: Arc, + } + + impl AsRef<[u8]> for DropTrackingOwner { + fn as_ref(&self) -> &[u8] { + &self.data + } + } + + impl Drop for DropTrackingOwner { + fn drop(&mut self) { + self.dropped.store(true, Ordering::Relaxed); + } + } + + #[tokio::test] + async fn uploads_full_part_before_finish() -> Result<()> { + let part_sizes = Arc::new(Mutex::new(Vec::new())); + let upload = RecordingMultipartUpload { + part_sizes: Arc::clone(&part_sizes), + }; + let mut writer = ObjectStoreSpillWriter { + upload: Box::new(upload), + buffer: PutPayloadMut::new(), + size: Arc::new(AtomicU64::new(0)), + bytes_written: 0, + }; - self.size.store(size, Ordering::Relaxed); + writer + .write_all(Bytes::from(vec![42; OBJECT_STORE_PART_SIZE + 1])) + .await?; + assert_eq!( + part_sizes.lock().unwrap().as_slice(), + &[OBJECT_STORE_PART_SIZE] + ); + assert_eq!(writer.buffer.content_length(), 1); + writer.abort().await?; Ok(()) } -} -/// Run an async ObjectStore operation. -/// -/// Adding a native async API is tracked in -fn block_on_object_store(future: impl Future>) -> Result { - if let Ok(handle) = tokio::runtime::Handle::try_current() { - tokio::task::block_in_place(|| handle.block_on(future)) - } else { - exec_err!("No current Tokio runtime available") + #[tokio::test] + async fn buffered_tail_does_not_retain_input_allocation() -> Result<()> { + let tmp_dir = tempdir()?; + let store: Arc = + Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path())?); + let location = Path::from("tail-retention-test"); + let upload = store.put_multipart(&location).await?; + let dropped = Arc::new(AtomicBool::new(false)); + let data = Bytes::from_owner(DropTrackingOwner { + data: vec![42; 128], + dropped: Arc::clone(&dropped), + }); + let mut writer = ObjectStoreSpillWriter { + upload, + buffer: PutPayloadMut::new(), + size: Arc::new(AtomicU64::new(0)), + bytes_written: 0, + }; + + writer.write_all(data).await?; + assert!( + dropped.load(Ordering::Relaxed), + "buffered tail should not retain the input allocation" + ); + writer.abort().await?; + Ok(()) } } diff --git a/datafusion/execution/src/lib.rs b/datafusion/execution/src/lib.rs index 5af7064f1cb8b..153c0216b84ad 100644 --- a/datafusion/execution/src/lib.rs +++ b/datafusion/execution/src/lib.rs @@ -49,6 +49,6 @@ pub mod registry { pub use async_stream::{Emitter, TryEmitter, async_stream, async_try_stream}; pub use disk_manager::DiskManager; pub use registry::FunctionRegistry; -pub use spill_file::{SpillFile, SpillWriter, TempFileFactory}; +pub use spill_file::{AsyncSpillWriter, SpillFile, SpillWriter, TempFileFactory}; pub use stream::{RecordBatchStream, SendableRecordBatchStream}; pub use task::{TaskContext, TaskContextProvider}; diff --git a/datafusion/execution/src/spill_file.rs b/datafusion/execution/src/spill_file.rs index dca5da23f53e1..d63d9e622e2eb 100644 --- a/datafusion/execution/src/spill_file.rs +++ b/datafusion/execution/src/spill_file.rs @@ -15,15 +15,18 @@ // specific language governing permissions and limitations // under the License. +use async_trait::async_trait; use bytes::Bytes; use datafusion_common::Result; use futures::Stream; +use std::io::Write; use std::path::Path; use std::pin::Pin; use std::sync::Arc; /// Abstraction over a spill file backend. /// Implementations handle their own quota enforcement and blocking concerns. +#[async_trait] pub trait SpillFile: Send + Sync { /// Returns the OS path if this is a local file, None otherwise. fn path(&self) -> Option<&Path> { @@ -38,14 +41,89 @@ pub trait SpillFile: Send + Sync { /// Opens a writer for appending data to this file. fn open_writer(&self) -> Result>; + + /// Opens an asynchronous writer for appending data to this file. + /// + /// The default implementation adapts [`Self::open_writer`] for backwards + /// compatibility. Backends with native asynchronous I/O should override + /// this method to avoid blocking an async executor. + async fn open_async_writer(&self) -> Result> { + Ok(Box::new(BlockingSpillWriterAdapter { + inner: self.open_writer()?, + })) + } } /// Writer for spill file backends. -pub trait SpillWriter: std::io::Write + Send { +pub trait SpillWriter: Write + Send { /// Intended for close/sync/commit operations. fn finish(&mut self) -> Result<()>; } +/// Asynchronous writer for spill file backends. +/// +/// The writer accepts owned [`Bytes`] so asynchronous backends can retain a +/// buffer across an await within [`Self::write_all`] without copying it. Calls +/// are made sequentially and buffers must be persisted in the order received. +/// After `write_all` returns, implementations must not retain the input's +/// backing allocation unless that retained memory is accounted for separately. +/// Backends that buffer a tail for a later call should copy that tail into an +/// allocation sized for the retained bytes. +/// +/// DataFusion makes a best-effort attempt to abort writers dropped by query +/// cancellation. Backends should still configure lifecycle cleanup for +/// abandoned uploads because cleanup cannot run after process termination. +#[async_trait] +pub trait AsyncSpillWriter: Send { + /// Writes all bytes in `data` to the spill file. + async fn write_all(&mut self, data: Bytes) -> Result<()>; + + /// Flushes buffered data, if supported by the backend. + /// + /// This does not finish the writer or require the spill file to become + /// visible to readers. [`Self::finish`] is the commit boundary. + async fn flush(&mut self) -> Result<()> { + Ok(()) + } + + /// Finishes and commits the spill file. + /// + /// A successful call is terminal. If committing returns an error, the + /// caller will invoke [`Self::abort`] before dropping the writer. + async fn finish(&mut self) -> Result<()>; + + /// Aborts an uncommitted spill write and cleans up backend resources. + /// + /// Multipart backends should override this method to explicitly abort an + /// in-progress upload. A successful call is terminal. DataFusion may retry + /// an abort that returns an error, including during best-effort drop cleanup. + async fn abort(&mut self) -> Result<()> { + Ok(()) + } +} + +struct BlockingSpillWriterAdapter { + inner: Box, +} + +#[async_trait] +impl AsyncSpillWriter for BlockingSpillWriterAdapter { + async fn write_all(&mut self, data: Bytes) -> Result<()> { + self.inner.write_all(&data)?; + Ok(()) + } + + async fn flush(&mut self) -> Result<()> { + self.inner.flush()?; + Ok(()) + } + + async fn finish(&mut self) -> Result<()> { + self.inner.flush()?; + self.inner.finish() + } +} + /// Factory for creating spill files. pub trait TempFileFactory: Send + Sync { fn create_temp_file(&self, description: &str) -> Result>; diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 490ea7cc85776..1cbd391fd7534 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -421,7 +421,7 @@ impl ExternalSorter { /// Appending globally sorted batches to the in-progress spill file, and clears /// the `globally_sorted_batches` (also its memory reservation) afterwards. - fn consume_and_spill_append( + async fn consume_and_spill_append( &mut self, globally_sorted_batches: &mut Vec, ) -> Result<()> { @@ -438,7 +438,9 @@ impl ExternalSorter { debug!("Spilling sort data of ExternalSorter to disk whilst inserting"); let batches_to_spill = std::mem::take(globally_sorted_batches); - self.reservation.free(); + // Keep the reservation alive while the batches remain in memory across + // asynchronous writes. It is released on success or error via RAII. + let _spill_reservation = self.reservation.take(); let (in_progress_file, max_record_batch_size) = self.in_progress_spill_file.as_mut().ok_or_else(|| { @@ -446,7 +448,7 @@ impl ExternalSorter { })?; for batch in batches_to_spill { - let gc_sliced_size = in_progress_file.append_batch(&batch)?; + let gc_sliced_size = in_progress_file.append_batch_async(&batch).await?; *max_record_batch_size = (*max_record_batch_size).max(gc_sliced_size); } @@ -460,12 +462,12 @@ impl ExternalSorter { } /// Finishes the in-progress spill file and moves it to the finished spill files. - fn spill_finish(&mut self) -> Result<()> { + async fn spill_finish(&mut self) -> Result<()> { let (mut in_progress_file, max_record_batch_memory) = self.in_progress_spill_file.take().ok_or_else(|| { internal_datafusion_err!("Should be called after `spill_append`") })?; - let spill_file = in_progress_file.finish()?; + let spill_file = in_progress_file.finish_async().await?; if let Some(spill_file) = spill_file { self.finished_spill_files.push(SortedSpillFile { @@ -477,9 +479,26 @@ impl ExternalSorter { Ok(()) } + async fn abort_in_progress_spill(&mut self) { + if let Some((in_progress_file, _)) = &mut self.in_progress_spill_file + && let Err(error) = in_progress_file.abort_async().await + { + debug!("Failed to abort in-progress sort spill: {error}"); + } + self.in_progress_spill_file.take(); + } + /// Sorts the in-memory batches and merges them into a single sorted run, then writes /// the result to spill files. async fn sort_and_spill_in_mem_batches(&mut self) -> Result<()> { + let result = self.try_sort_and_spill_in_mem_batches().await; + if result.is_err() { + self.abort_in_progress_spill().await; + } + result + } + + async fn try_sort_and_spill_in_mem_batches(&mut self) -> Result<()> { assert_or_internal_err!( !self.in_mem_batches.is_empty(), "in_mem_batches must not be empty when attempting to sort and spill" @@ -509,13 +528,22 @@ impl ExternalSorter { while let Some(batch) = sorted_stream.next().await { let batch = batch?; let sorted_size = get_reserved_bytes_for_record_batch(&batch)?; - let reservation_failed = self.reservation.try_grow(sorted_size).is_err(); + let reservation_failed = match self.reservation.try_grow(sorted_size) { + Ok(()) => false, + Err(_) => { + // The batch is already materialized, so account for it while + // spilling even if this temporarily exceeds the pool limit. + self.reservation.grow(sorted_size); + true + } + }; // Even if the reservation is not enough, the batch is already in // memory, so it's okay to combine it with previously sorted // batches, and spill together. globally_sorted_batches.push(batch); if reservation_failed { - self.consume_and_spill_append(&mut globally_sorted_batches)?; // reservation is freed in spill() + self.consume_and_spill_append(&mut globally_sorted_batches) + .await?; // reservation is released when the spill completes } } @@ -523,8 +551,9 @@ impl ExternalSorter { // upcoming `self.reserve_memory_for_merge()` may fail due to insufficient memory. drop(sorted_stream); - self.consume_and_spill_append(&mut globally_sorted_batches)?; - self.spill_finish()?; + self.consume_and_spill_append(&mut globally_sorted_batches) + .await?; + self.spill_finish().await?; // Sanity check after spilling let buffers_cleared_property = @@ -1490,8 +1519,17 @@ impl ExecutionPlan for SortExec { self.schema(), futures::stream::once(async move { while let Some(batch) = input.next().await { - let batch = batch?; - sorter.insert_batch(batch).await?; + let batch = match batch { + Ok(batch) => batch, + Err(error) => { + sorter.abort_in_progress_spill().await; + return Err(error); + } + }; + if let Err(error) = sorter.insert_batch(batch).await { + sorter.abort_in_progress_spill().await; + return Err(error); + } } drop(input); sorter.sort().await @@ -2115,7 +2153,9 @@ mod proto_tests { #[cfg(test)] mod tests { use std::collections::HashMap; + use std::path::Path; use std::pin::Pin; + use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::task::{Context, Poll}; use super::*; @@ -2133,22 +2173,106 @@ mod tests { use arrow::array::*; use arrow::compute::SortOptions; use arrow::datatypes::*; + use async_trait::async_trait; + use bytes::Bytes; use datafusion_common::ScalarValue; use datafusion_common::cast::as_primitive_array; use datafusion_common::config::ConfigOptions; use datafusion_common::test_util::batches_to_string; use datafusion_execution::RecordBatchStream; use datafusion_execution::config::SessionConfig; + use datafusion_execution::disk_manager::DiskManagerBuilder; use datafusion_execution::memory_pool::{ GreedyMemoryPool, MemoryConsumer, MemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_execution::{ + AsyncSpillWriter, SpillFile, SpillWriter, TempFileFactory, + }; use datafusion_physical_expr::expressions::{Column, Literal}; use datafusion_physical_expr::{DynamicFilterTracking, EquivalenceProperties}; use datafusion_physical_expr_common::metrics::MetricValue; use futures::{FutureExt, Stream, TryStreamExt}; use insta::assert_snapshot; + use tokio::sync::Notify; + + struct PendingWriteTempFileFactory { + write_started: Arc, + aborted: Arc, + abort_count: Arc, + } + + impl TempFileFactory for PendingWriteTempFileFactory { + fn create_temp_file(&self, _description: &str) -> Result> { + Ok(Arc::new(PendingWriteSpillFile { + write_started: Arc::clone(&self.write_started), + aborted: Arc::clone(&self.aborted), + abort_count: Arc::clone(&self.abort_count), + })) + } + } + + struct PendingWriteSpillFile { + write_started: Arc, + aborted: Arc, + abort_count: Arc, + } + + #[async_trait] + impl SpillFile for PendingWriteSpillFile { + fn path(&self) -> Option<&Path> { + None + } + + fn size(&self) -> Option { + Some(0) + } + + fn read_stream( + &self, + ) -> Result> + Send>>> { + Ok(Box::pin(futures::stream::empty())) + } + + fn open_writer(&self) -> Result> { + datafusion_common::not_impl_err!( + "test backend only supports asynchronous writes" + ) + } + + async fn open_async_writer(&self) -> Result> { + Ok(Box::new(PendingWriteWriter { + write_started: Arc::clone(&self.write_started), + aborted: Arc::clone(&self.aborted), + abort_count: Arc::clone(&self.abort_count), + })) + } + } + + struct PendingWriteWriter { + write_started: Arc, + aborted: Arc, + abort_count: Arc, + } + + #[async_trait] + impl AsyncSpillWriter for PendingWriteWriter { + async fn write_all(&mut self, _data: Bytes) -> Result<()> { + self.write_started.notify_one(); + std::future::pending().await + } + + async fn finish(&mut self) -> Result<()> { + Ok(()) + } + + async fn abort(&mut self) -> Result<()> { + self.abort_count.fetch_add(1, AtomicOrdering::Relaxed); + self.aborted.notify_one(); + Ok(()) + } + } #[derive(Debug, Clone)] pub struct SortedUnboundedExec { @@ -3831,6 +3955,75 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_spill_reservation_held_during_async_write() -> Result<()> { + let pool: Arc = Arc::new(GreedyMemoryPool::new(0)); + let write_started = Arc::new(Notify::new()); + let aborted = Arc::new(Notify::new()); + let abort_count = Arc::new(AtomicUsize::new(0)); + let disk_manager_builder = DiskManagerBuilder::default().with_temp_file_factory( + Arc::new(PendingWriteTempFileFactory { + write_started: Arc::clone(&write_started), + aborted: Arc::clone(&aborted), + abort_count: Arc::clone(&abort_count), + }), + ); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .with_disk_manager_builder(disk_manager_builder) + .build_arc()?; + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); + let metrics = ExecutionPlanMetricsSet::new(); + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(), + 128, + 0, + usize::MAX, + SpillCompression::Uncompressed, + &metrics, + runtime, + )?; + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![3, 2, 1]))], + )?; + let reserved_bytes = get_reserved_bytes_for_record_batch(&batch)?; + sorter.reservation.grow(reserved_bytes); + sorter.in_mem_batches.push(batch); + + #[expect(clippy::disallowed_methods)] // spawn allowed only in tests + let task = + tokio::spawn(async move { sorter.sort_and_spill_in_mem_batches().await }); + + if tokio::time::timeout( + std::time::Duration::from_secs(5), + write_started.notified(), + ) + .await + .is_err() + { + task.abort(); + let _ = task.await; + panic!("spill write did not start before the timeout"); + } + assert_eq!( + pool.reserved(), + reserved_bytes, + "resident batches must remain accounted for while an async spill is pending" + ); + + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + tokio::time::timeout(std::time::Duration::from_secs(1), aborted.notified()) + .await + .expect("cancelling the spill should abort its writer"); + assert_eq!(abort_count.load(AtomicOrdering::Relaxed), 1); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + /// Verifies that `ExternalSorter::sort()` transfers the pre-reserved /// merge bytes to the merge stream via `take()`, rather than leaving /// them in the sorter (via `new_empty()`). diff --git a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs index 71d7cce1bcc7d..2f4ddd3398f8e 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -17,7 +17,7 @@ //! Define the `InProgressSpillFile` struct, which represents an in-progress spill file used for writing `RecordBatch`es to disk, created by `SpillManager`. -use datafusion_common::Result; +use datafusion_common::{Result, internal_datafusion_err}; use std::sync::Arc; use arrow::array::RecordBatch; @@ -25,28 +25,33 @@ use datafusion_common::exec_datafusion_err; use datafusion_execution::spill_file::SpillFile; use super::{ - IPCStreamWriter, gc_view_arrays, + AsyncIPCStreamWriter, IPCStreamEncoder, IPCStreamWriter, gc_view_arrays, spill_manager::{GetSlicedSize, SpillManager}, }; +enum InProgressWriter { + Sync(IPCStreamWriter), + Async(AsyncIPCStreamWriter), +} + /// Represents an in-progress spill file used for writing `RecordBatch`es to disk, created by `SpillManager`. /// Caller is able to use this struct to incrementally append in-memory batches to /// the file, and then finalize the file by calling the `finish` method. pub struct InProgressSpillFile { - pub(crate) spill_writer: Arc, + pub(crate) spill_manager: Arc, /// Lazily initialized writer - writer: Option, + writer: Option, /// Lazily initialized in-progress file, it will be moved out when the `finish` method is invoked in_progress_file: Option>, } impl InProgressSpillFile { pub fn new( - spill_writer: Arc, + spill_manager: Arc, in_progress_file: Arc, ) -> Self { Self { - spill_writer, + spill_manager, in_progress_file: Some(in_progress_file), writer: None, } @@ -78,35 +83,92 @@ impl InProgressSpillFile { // Individual batches may have different schemas (e.g., different nullability) // when they come from different branches of a UnionExec. The SpillManager's // schema represents the canonical schema that all batches should conform to. - let schema = self.spill_writer.schema(); + let schema = self.spill_manager.schema(); if let Some(in_progress_file) = &self.in_progress_file { let spill_writer = in_progress_file.open_writer()?; - self.writer = Some(IPCStreamWriter::new( + self.writer = Some(InProgressWriter::Sync(IPCStreamWriter::new( spill_writer, schema.as_ref(), - self.spill_writer.compression, - )?); + self.spill_manager.compression, + )?)); // Update metrics - self.spill_writer.metrics.spill_file_count.add(1); - let header_bytes = self.writer.as_ref().unwrap().bytes_written(); - self.spill_writer.metrics.spilled_bytes.add(header_bytes); + self.spill_manager.metrics.spill_file_count.add(1); } } - if let Some(writer) = &mut self.writer { + if let Some(InProgressWriter::Sync(writer)) = &mut self.writer { // The writer calculates how many serialized bytes were emitted let (spilled_rows, delta_bytes) = writer.write(&gc_batch)?; - self.spill_writer.metrics.spilled_rows.add(spilled_rows); - self.spill_writer.metrics.spilled_bytes.add(delta_bytes); + self.spill_manager.metrics.spilled_rows.add(spilled_rows); + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + } else if self.writer.is_some() { + return Err(exec_datafusion_err!( + "Cannot use synchronous append after asynchronous spill writing has started" + )); + } + gc_batch.get_sliced_size() + } + + /// Appends a `RecordBatch` using the asynchronous spill writer. + pub async fn append_batch_async(&mut self, batch: &RecordBatch) -> Result { + if self.in_progress_file.is_none() { + return Err(exec_datafusion_err!( + "Append operation failed: No active in-progress file. The file may have already been finalized." + )); + } + + let gc_batch = gc_view_arrays(batch)?; + + if self.writer.is_none() { + let schema = self.spill_manager.schema(); + if let Some(in_progress_file) = &self.in_progress_file { + // Validate the IPC schema and compression options before opening + // a remote upload that would otherwise need to be aborted. + let encoder = IPCStreamEncoder::new( + schema.as_ref(), + self.spill_manager.compression, + )?; + let spill_writer = in_progress_file.open_async_writer().await?; + + self.writer = Some(InProgressWriter::Async(AsyncIPCStreamWriter::new( + spill_writer, + encoder, + ))); + + self.spill_manager.metrics.spill_file_count.add(1); + } + } + + match &mut self.writer { + Some(InProgressWriter::Async(writer)) => { + let (spilled_rows, delta_bytes) = writer.write(&gc_batch).await?; + + self.spill_manager.metrics.spilled_rows.add(spilled_rows); + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + } + Some(InProgressWriter::Sync(_)) => { + return Err(exec_datafusion_err!( + "Cannot use asynchronous append after synchronous spill writing has started" + )); + } + None => { + return Err(internal_datafusion_err!( + "Asynchronous spill writer was not initialized" + )); + } } gc_batch.get_sliced_size() } pub fn flush(&mut self) -> Result<()> { - if let Some(writer) = &mut self.writer { + if let Some(InProgressWriter::Sync(writer)) = &mut self.writer { writer.flush()?; + } else if self.writer.is_some() { + return Err(exec_datafusion_err!( + "Cannot use synchronous flush for an asynchronous spill writer" + )); } Ok(()) } @@ -125,16 +187,62 @@ impl InProgressSpillFile { "Finish operation failed: file has already been finalized." )); } - if let Some(mut writer) = self.writer.take() { + if matches!(self.writer, Some(InProgressWriter::Async(_))) { + return Err(exec_datafusion_err!( + "Cannot use synchronous finish for an asynchronous spill writer" + )); + } + if let Some(InProgressWriter::Sync(mut writer)) = self.writer.take() { // Finish the writer and capture any final trailing bytes emitted let delta_bytes = writer.finish()?; - self.spill_writer.metrics.spilled_bytes.add(delta_bytes); + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + } else { + return Ok(None); + } + + Ok(self.in_progress_file.take()) + } + + /// Finalizes an asynchronous spill write, returning the completed file. + pub async fn finish_async(&mut self) -> Result>> { + if self.in_progress_file.is_none() && self.writer.is_none() { + return Err(exec_datafusion_err!( + "Finish operation failed: file has already been finalized." + )); + } + if matches!(self.writer, Some(InProgressWriter::Sync(_))) { + return Err(exec_datafusion_err!( + "Cannot use asynchronous finish for a synchronous spill writer" + )); + } + if let Some(InProgressWriter::Async(writer)) = &mut self.writer { + let delta_bytes = writer.finish().await?; + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); } else { return Ok(None); } + self.writer.take(); Ok(self.in_progress_file.take()) } + + /// Aborts an asynchronous spill write and discards its file. + pub async fn abort_async(&mut self) -> Result<()> { + if matches!(self.writer, Some(InProgressWriter::Sync(_))) { + return Err(exec_datafusion_err!( + "Cannot use asynchronous abort for a synchronous spill writer" + )); + } + + let result = if let Some(InProgressWriter::Async(writer)) = &mut self.writer { + writer.abort().await + } else { + Ok(()) + }; + self.writer.take(); + self.in_progress_file.take(); + result + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index e3c4b2cf30f9c..49197326e4917 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -21,7 +21,7 @@ pub(crate) mod in_progress_spill_file; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; pub mod spill_pool; -use datafusion_execution::spill_file::SpillWriter; +use datafusion_execution::spill_file::{AsyncSpillWriter, SpillWriter}; // Moved for refactor, re-export to keep the public API stable pub use datafusion_common::utils::memory::get_record_batch_memory_size; // Re-export SpillManager for doctests only (hidden from public docs) @@ -30,8 +30,9 @@ pub use spill_manager::SpillManager; use std::collections::VecDeque; use std::pin::Pin; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; +use std::time::Duration; use arrow::array::{ Array, ArrayRef, BinaryViewArray, BufferSpec, GenericByteViewArray, StringViewArray, @@ -43,14 +44,16 @@ use arrow::datatypes::{ByteViewType, Schema, SchemaRef}; use arrow::ipc::{ MetadataVersion, reader::StreamDecoder, - writer::{IpcWriteOptions, StreamWriter}, + writer::{IpcWriteOptions, StreamEncoder}, }; use arrow::record_batch::RecordBatch; use arrow_data::ArrayDataBuilder; +#[cfg(test)] +use arrow_ipc::writer::StreamWriter; use arrow_ipc::{CompressionType, root_as_message}; -use datafusion_common::Result; use datafusion_common::config::SpillCompression; +use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::RecordBatchStream; use datafusion_execution::spill_file::SpillFile; use futures::Stream; @@ -464,59 +467,210 @@ impl RecordBatchStream for SpillReaderStream { } } -/// A wrapper that counts the exact compressed IPC bytes written by Arrow. -/// -/// Arrow's `StreamWriter` does not return the number of bytes written during its -/// `write()` calls. To accurately track the `spilled_bytes` metrics (especially -/// when LZ4/ZSTD compression is applied), we must intercept the `std::io::Write` -/// trait boundary to count the final serialized payload size. -pub(crate) struct TrackingSpillWriter { - inner: Box, - pub(crate) total_bytes_written: usize, +/// Write in Arrow IPC Stream format to an underlying `SpillWriter` backend. +/// Stream format also supports dictionary replacement. +struct IPCStreamWriter { + writer: Box, + encoder: IPCStreamEncoder, } -impl TrackingSpillWriter { - pub fn new(inner: Box) -> Self { - Self { - inner, - total_bytes_written: 0, +impl IPCStreamWriter { + pub fn new( + spill_writer: Box, + schema: &Schema, + spill_compression: SpillCompression, + ) -> Result { + Ok(Self { + writer: spill_writer, + encoder: IPCStreamEncoder::new(schema, spill_compression)?, + }) + } + + pub fn write(&mut self, batch: &RecordBatch) -> Result<(usize, usize)> { + use std::io::Write; + let (rows, bytes, buffers) = self.encoder.encode(batch)?; + for buffer in buffers { + self.writer.write_all(buffer.as_slice())?; } + Ok((rows, bytes)) + } + + pub fn flush(&mut self) -> Result<()> { + use std::io::Write; + self.writer.flush()?; + Ok(()) } - pub fn finish(mut self) -> Result<()> { - self.inner.finish() + pub fn finish(&mut self) -> Result { + use std::io::Write; + let (bytes, buffers) = self.encoder.finish()?; + for buffer in buffers { + self.writer.write_all(buffer.as_slice())?; + } + self.writer.flush()?; + self.writer.finish()?; + Ok(bytes) } } -impl std::io::Write for TrackingSpillWriter { - fn write(&mut self, buf: &[u8]) -> std::io::Result { - let n = self.inner.write(buf)?; +/// Writes Arrow IPC Stream data through an [`AsyncSpillWriter`] interface. +struct AsyncIPCStreamWriter { + writer: Option>, + encoder: IPCStreamEncoder, + abort_handle: Option, +} - self.total_bytes_written += n; +/// Bounds detached cleanup work after query cancellation. If all permits are +/// occupied, the writer is dropped and the backend lifecycle policy is the +/// remaining cleanup mechanism. +const MAX_CONCURRENT_SPILL_ABORTS: usize = 8; +const SPILL_ABORT_TIMEOUT: Duration = Duration::from_secs(30); +static SPILL_ABORT_PERMITS: LazyLock> = + LazyLock::new(|| Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_SPILL_ABORTS))); - Ok(n) +impl AsyncIPCStreamWriter { + pub fn new( + spill_writer: Box, + encoder: IPCStreamEncoder, + ) -> Self { + Self { + writer: Some(spill_writer), + encoder, + abort_handle: tokio::runtime::Handle::try_current().ok(), + } } - fn flush(&mut self) -> std::io::Result<()> { - self.inner.flush() + pub async fn write(&mut self, batch: &RecordBatch) -> Result<(usize, usize)> { + let (rows, bytes, buffers) = match self.encoder.encode(batch) { + Ok(encoded) => encoded, + Err(error) => { + self.abort_after_error().await; + return Err(error); + } + }; + for buffer in buffers { + let result = self + .writer + .as_mut() + .ok_or_else(|| { + internal_datafusion_err!("Spill writer is no longer active") + })? + .write_all(buffer_to_bytes(buffer)) + .await; + if let Err(error) = result { + self.abort_after_error().await; + return Err(error); + } + } + Ok((rows, bytes)) + } + + pub async fn finish(&mut self) -> Result { + let (bytes, buffers) = match self.encoder.finish() { + Ok(encoded) => encoded, + Err(error) => { + self.abort_after_error().await; + return Err(error); + } + }; + for buffer in buffers { + let result = self + .writer + .as_mut() + .ok_or_else(|| { + internal_datafusion_err!("Spill writer is no longer active") + })? + .write_all(buffer_to_bytes(buffer)) + .await; + if let Err(error) = result { + self.abort_after_error().await; + return Err(error); + } + } + let result = self + .writer + .as_mut() + .ok_or_else(|| internal_datafusion_err!("Spill writer is no longer active"))? + .flush() + .await; + if let Err(error) = result { + self.abort_after_error().await; + return Err(error); + } + let result = self + .writer + .as_mut() + .ok_or_else(|| internal_datafusion_err!("Spill writer is no longer active"))? + .finish() + .await; + if let Err(error) = result { + self.abort_after_error().await; + return Err(error); + } + self.writer.take(); + Ok(bytes) + } + + pub async fn abort(&mut self) -> Result<()> { + let Some(writer) = self.writer.as_mut() else { + return Ok(()); + }; + writer.abort().await?; + self.writer = None; + Ok(()) + } + + async fn abort_after_error(&mut self) { + if let Err(error) = self.abort().await { + debug!("Failed to abort spill writer: {error}"); + } } } -/// Write in Arrow IPC Stream format to an underlying `SpillWriter` backend. -/// Stream format also supports dictionary replacement. -struct IPCStreamWriter { - /// Inner writer - writer: Option>, - /// Batches written - num_batches: usize, - /// Rows written - num_rows: usize, - /// Bytes written - num_bytes: usize, +impl Drop for AsyncIPCStreamWriter { + fn drop(&mut self) { + let Some(mut writer) = self.writer.take() else { + return; + }; + + let Some(handle) = tokio::runtime::Handle::try_current() + .ok() + .or_else(|| self.abort_handle.take()) + else { + debug!("Unable to abort dropped spill writer without a Tokio runtime"); + return; + }; + + let Ok(permit) = Arc::clone(&SPILL_ABORT_PERMITS).try_acquire_owned() else { + debug!( + "Skipping spill writer abort because the cleanup concurrency limit was reached" + ); + return; + }; + + let _abort_task = handle.spawn(async move { + match tokio::time::timeout(SPILL_ABORT_TIMEOUT, writer.abort()).await { + Ok(Ok(())) => {} + Ok(Err(error)) => { + debug!("Failed to abort dropped spill writer: {error}"); + } + Err(_) => { + debug!("Timed out aborting dropped spill writer"); + } + } + // `permit` is released when this bounded cleanup task exits. + drop(permit); + }); + } } -impl IPCStreamWriter { - /// Create new writer +/// Encodes Arrow IPC streams without performing I/O. +struct IPCStreamEncoder { + encoder: Option, +} + +impl IPCStreamEncoder { + /// Create a new encoder. /// /// # Codec contract /// @@ -528,11 +682,7 @@ impl IPCStreamWriter { /// contract local and build-visible during Cargo feature resolution, /// rather than relying solely on workspace-level feature unification; /// see #21917. - pub fn new( - spill_writer: Box, - schema: &Schema, - spill_compression: SpillCompression, - ) -> Result { + fn new(schema: &Schema, spill_compression: SpillCompression) -> Result { let metadata_version = MetadataVersion::V5; // Depending on the schema, some array types such as StringViewArray require larger (16 byte in this case) alignment. // If the actual buffer layout after IPC read does not satisfy the alignment requirement, @@ -546,68 +696,37 @@ impl IPCStreamWriter { let compression_type = Option::::from(spill_compression); write_options = write_options.try_with_compression(compression_type)?; - let adapter = TrackingSpillWriter::new(spill_writer); - let writer = StreamWriter::try_new_with_options(adapter, schema, write_options)?; - + let encoder = StreamEncoder::try_new_with_options(schema, write_options)?; Ok(Self { - num_batches: 0, - num_rows: 0, - num_bytes: 0, - writer: Some(writer), + encoder: Some(encoder), }) } - /// Writes a single batch to the IPC stream and updates the internal counters. - /// - /// Returns a tuple containing the change in the number of rows and bytes written. - pub fn write(&mut self, batch: &RecordBatch) -> Result<(usize, usize)> { - let writer = self.writer.as_mut().unwrap(); - - let bytes_before = writer.get_ref().total_bytes_written; - writer.write(batch)?; - let bytes_after = writer.get_ref().total_bytes_written; - self.num_batches += 1; - let delta_num_rows = batch.num_rows(); - self.num_rows += delta_num_rows; - let delta_num_bytes = bytes_after - bytes_before; - self.num_bytes += delta_num_bytes; - Ok((delta_num_rows, delta_num_bytes)) + fn encode(&mut self, batch: &RecordBatch) -> Result<(usize, usize, Vec)> { + let buffers = self.encoder.as_mut().unwrap().encode(batch)?; + let bytes = buffers.iter().map(Buffer::len).sum(); + Ok((batch.num_rows(), bytes, buffers)) } - pub fn flush(&mut self) -> Result<()> { - use std::io::Write; - if let Some(writer) = &mut self.writer { - writer.get_mut().flush()?; - } - Ok(()) + fn finish(&mut self) -> Result<(usize, Vec)> { + let buffers = self.encoder.take().unwrap().finish()?; + let bytes = buffers.iter().map(Buffer::len).sum(); + Ok((bytes, buffers)) } +} - /// Finish the writer. - /// - /// Returns the number of trailing bytes written during the finish operation - /// (e.g., IPC metadata and footers). - pub fn finish(&mut self) -> Result { - let mut writer = self.writer.take().unwrap(); - - let bytes_before = writer.get_ref().total_bytes_written; - writer.finish()?; // Writes IPC tail - - // Extract the adapter and flush the final bytes - let adapter = writer.into_inner()?; - let bytes_after = adapter.total_bytes_written; - adapter.finish()?; +struct ArrowBufferOwner(Buffer); - Ok(bytes_after - bytes_before) - } - /// Returns the total number of bytes written so far - pub fn bytes_written(&self) -> usize { - self.writer - .as_ref() - .map(|w| w.get_ref().total_bytes_written) - .unwrap_or(0) +impl AsRef<[u8]> for ArrowBufferOwner { + fn as_ref(&self) -> &[u8] { + self.0.as_slice() } } +fn buffer_to_bytes(buffer: Buffer) -> bytes::Bytes { + bytes::Bytes::from_owner(ArrowBufferOwner(buffer)) +} + // Returns the maximum byte alignment required by any field in the schema (>= 8), derived from Arrow buffer layouts. fn get_max_alignment_for_schema(schema: &Schema) -> usize { let minimum_alignment = 8; @@ -811,8 +930,308 @@ mod tests { use arrow::array::{ArrayRef, Int32Array, StringArray}; use arrow::compute::cast; use arrow::datatypes::{DataType, Field}; + use arrow::ipc::reader::StreamReader; use datafusion_execution::runtime_env::RuntimeEnv; use futures::StreamExt as _; + use std::io::Cursor; + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use tokio::sync::Notify; + + struct YieldingAsyncWriter { + data: Arc>>, + writes: Arc, + finished: Arc, + aborted: Arc, + } + + struct FailingAsyncWriter { + aborted: Arc, + } + + struct PendingFinishAsyncWriter { + finish_started: Arc, + aborted: Arc, + abort_count: Arc, + } + + struct FailOnceAbortWriter { + abort_count: Arc, + cleanup_done: Arc, + } + + #[async_trait::async_trait] + impl AsyncSpillWriter for FailingAsyncWriter { + async fn write_all(&mut self, _data: bytes::Bytes) -> Result<()> { + datafusion_common::exec_err!("injected spill write failure") + } + + async fn finish(&mut self) -> Result<()> { + Ok(()) + } + + async fn abort(&mut self) -> Result<()> { + self.aborted.store(true, Ordering::Relaxed); + Ok(()) + } + } + + #[async_trait::async_trait] + impl AsyncSpillWriter for YieldingAsyncWriter { + async fn write_all(&mut self, data: bytes::Bytes) -> Result<()> { + tokio::task::yield_now().await; + self.writes.fetch_add(1, Ordering::Relaxed); + self.data.lock().unwrap().extend_from_slice(&data); + Ok(()) + } + + async fn finish(&mut self) -> Result<()> { + self.finished.store(true, Ordering::Relaxed); + Ok(()) + } + + async fn abort(&mut self) -> Result<()> { + self.aborted.store(true, Ordering::Relaxed); + Ok(()) + } + } + + #[async_trait::async_trait] + impl AsyncSpillWriter for PendingFinishAsyncWriter { + async fn write_all(&mut self, _data: bytes::Bytes) -> Result<()> { + Ok(()) + } + + async fn finish(&mut self) -> Result<()> { + self.finish_started.notify_one(); + std::future::pending().await + } + + async fn abort(&mut self) -> Result<()> { + self.abort_count.fetch_add(1, Ordering::Relaxed); + self.aborted.notify_one(); + Ok(()) + } + } + + #[async_trait::async_trait] + impl AsyncSpillWriter for FailOnceAbortWriter { + async fn write_all(&mut self, _data: bytes::Bytes) -> Result<()> { + Ok(()) + } + + async fn finish(&mut self) -> Result<()> { + Ok(()) + } + + async fn abort(&mut self) -> Result<()> { + if self.abort_count.fetch_add(1, Ordering::Relaxed) == 0 { + return datafusion_common::exec_err!("injected spill abort failure"); + } + self.cleanup_done.notify_one(); + Ok(()) + } + } + + #[tokio::test] + async fn test_async_ipc_stream_writer() -> Result<()> { + let batch1 = build_table_i32( + ("a", &vec![0, 1, 2]), + ("b", &vec![3, 4, 5]), + ("c", &vec![6, 7, 8]), + ); + let batch2 = build_table_i32( + ("a", &vec![9, 10, 11]), + ("b", &vec![12, 13, 14]), + ("c", &vec![15, 16, 17]), + ); + let schema = batch1.schema(); + let data = Arc::new(Mutex::new(Vec::new())); + let writes = Arc::new(AtomicUsize::new(0)); + let finished = Arc::new(AtomicBool::new(false)); + let aborted = Arc::new(AtomicBool::new(false)); + + let spill_writer = YieldingAsyncWriter { + data: Arc::clone(&data), + writes: Arc::clone(&writes), + finished: Arc::clone(&finished), + aborted: Arc::clone(&aborted), + }; + let encoder = + IPCStreamEncoder::new(schema.as_ref(), SpillCompression::Uncompressed)?; + let mut writer = AsyncIPCStreamWriter::new(Box::new(spill_writer), encoder); + + let (_, first_bytes) = writer.write(&batch1).await?; + let (_, second_bytes) = writer.write(&batch2).await?; + let trailing_bytes = writer.finish().await?; + + assert!(finished.load(Ordering::Relaxed)); + assert!(writes.load(Ordering::Relaxed) > 2); + drop(writer); + assert!(!aborted.load(Ordering::Relaxed)); + + let encoded = data.lock().unwrap().clone(); + assert_eq!(encoded.len(), first_bytes + second_bytes + trailing_bytes); + let decoded = StreamReader::try_new(Cursor::new(encoded), None)? + .collect::>>()?; + assert_eq!(decoded, vec![batch1, batch2]); + Ok(()) + } + + #[test] + fn test_buffer_to_bytes_is_zero_copy() { + let buffer = Buffer::from(vec![0_u8, 1, 2, 3]).slice_with_length(1, 2); + let expected_ptr = buffer.as_ptr(); + let bytes = buffer_to_bytes(buffer); + + assert_eq!(bytes.as_ptr(), expected_ptr); + assert_eq!(bytes.as_ref(), &[1, 2]); + } + + #[tokio::test] + async fn test_async_ipc_stream_writer_aborts_after_write_error() -> Result<()> { + let batch = build_table_i32( + ("a", &vec![0, 1, 2]), + ("b", &vec![3, 4, 5]), + ("c", &vec![6, 7, 8]), + ); + let aborted = Arc::new(AtomicBool::new(false)); + let encoder = IPCStreamEncoder::new( + batch.schema().as_ref(), + SpillCompression::Uncompressed, + )?; + let mut writer = AsyncIPCStreamWriter::new( + Box::new(FailingAsyncWriter { + aborted: Arc::clone(&aborted), + }), + encoder, + ); + + let error = writer.write(&batch).await.unwrap_err(); + assert!(error.to_string().contains("injected spill write failure")); + assert!(aborted.load(Ordering::Relaxed)); + Ok(()) + } + + #[tokio::test] + async fn test_async_ipc_stream_writer_aborts_cancelled_finish() -> Result<()> { + let batch = build_table_i32( + ("a", &vec![0, 1, 2]), + ("b", &vec![3, 4, 5]), + ("c", &vec![6, 7, 8]), + ); + let finish_started = Arc::new(Notify::new()); + let aborted = Arc::new(Notify::new()); + let abort_count = Arc::new(AtomicUsize::new(0)); + + #[expect(clippy::disallowed_methods)] // spawn allowed only in tests + let task = tokio::spawn({ + let finish_started = Arc::clone(&finish_started); + let aborted = Arc::clone(&aborted); + let abort_count = Arc::clone(&abort_count); + async move { + let encoder = IPCStreamEncoder::new( + batch.schema().as_ref(), + SpillCompression::Uncompressed, + )?; + let mut writer = AsyncIPCStreamWriter::new( + Box::new(PendingFinishAsyncWriter { + finish_started, + aborted, + abort_count, + }), + encoder, + ); + writer.write(&batch).await?; + writer.finish().await?; + Result::<()>::Ok(()) + } + }); + + if tokio::time::timeout(Duration::from_secs(5), finish_started.notified()) + .await + .is_err() + { + task.abort(); + let _ = task.await; + panic!("spill finish did not start before the timeout"); + } + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(1), aborted.notified()) + .await + .expect("dropping a cancelled writer should abort its upload"); + assert_eq!(abort_count.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[tokio::test] + async fn test_async_ipc_stream_writer_retries_failed_abort_on_drop() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let abort_count = Arc::new(AtomicUsize::new(0)); + let cleanup_done = Arc::new(Notify::new()); + let encoder = IPCStreamEncoder::new(&schema, SpillCompression::Uncompressed)?; + let mut writer = AsyncIPCStreamWriter::new( + Box::new(FailOnceAbortWriter { + abort_count: Arc::clone(&abort_count), + cleanup_done: Arc::clone(&cleanup_done), + }), + encoder, + ); + + assert!(writer.abort().await.is_err()); + assert_eq!(abort_count.load(Ordering::Relaxed), 1); + drop(writer); + + tokio::time::timeout(Duration::from_secs(1), cleanup_done.notified()) + .await + .expect("dropping a writer should retry an abort that returned an error"); + assert_eq!(abort_count.load(Ordering::Relaxed), 2); + Ok(()) + } + + #[test] + fn test_async_ipc_stream_writer_aborts_on_current_runtime() -> Result<()> { + let batch = build_table_i32( + ("a", &vec![0, 1, 2]), + ("b", &vec![3, 4, 5]), + ("c", &vec![6, 7, 8]), + ); + let aborted = Arc::new(AtomicBool::new(false)); + let first_runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + let writer = { + let _runtime_guard = first_runtime.enter(); + let spill_writer = YieldingAsyncWriter { + data: Arc::new(Mutex::new(Vec::new())), + writes: Arc::new(AtomicUsize::new(0)), + finished: Arc::new(AtomicBool::new(false)), + aborted: Arc::clone(&aborted), + }; + let encoder = IPCStreamEncoder::new( + batch.schema().as_ref(), + SpillCompression::Uncompressed, + )?; + AsyncIPCStreamWriter::new(Box::new(spill_writer), encoder) + }; + drop(first_runtime); + + let second_runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + second_runtime.block_on(async { + drop(writer); + tokio::time::timeout(Duration::from_secs(1), async { + while !aborted.load(Ordering::Relaxed) { + tokio::task::yield_now().await; + } + }) + .await + .expect("drop should abort on the current live runtime"); + }); + Ok(()) + } #[tokio::test] async fn test_batch_spill_and_read() -> Result<()> { @@ -1444,14 +1863,14 @@ mod tests { expected_spilled_rows: usize, ) -> Result<()> { let actual_spill_file_count = in_progress_file - .spill_writer + .spill_manager .metrics .spill_file_count .value(); let actual_spilled_bytes = - in_progress_file.spill_writer.metrics.spilled_bytes.value(); + in_progress_file.spill_manager.metrics.spilled_bytes.value(); let actual_spilled_rows = - in_progress_file.spill_writer.metrics.spilled_rows.value(); + in_progress_file.spill_manager.metrics.spilled_rows.value(); assert_eq!( actual_spill_file_count, expected_spill_file_count, diff --git a/datafusion/physical-plan/src/spill/spill_manager.rs b/datafusion/physical-plan/src/spill/spill_manager.rs index aee9e917c755d..1db4f539827a8 100644 --- a/datafusion/physical-plan/src/spill/spill_manager.rs +++ b/datafusion/physical-plan/src/spill/spill_manager.rs @@ -27,6 +27,7 @@ use datafusion_common::{DataFusionError, Result, config::SpillCompression}; use datafusion_execution::SendableRecordBatchStream; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_execution::spill_file::SpillFile; +use log::debug; use std::borrow::Borrow; use std::sync::Arc; @@ -150,18 +151,29 @@ impl SpillManager { let mut in_progress_file = self.create_in_progress_file(request_description)?; - let mut max_record_batch_size = 0; + let result = async { + let mut max_record_batch_size = 0; - while let Some(batch) = stream.next().await { - let batch = batch?; - let gc_sliced_size = in_progress_file.append_batch(&batch)?; + while let Some(batch) = stream.next().await { + let batch = batch?; + let gc_sliced_size = in_progress_file.append_batch_async(&batch).await?; - max_record_batch_size = max_record_batch_size.max(gc_sliced_size); + max_record_batch_size = max_record_batch_size.max(gc_sliced_size); + } + + let file = in_progress_file.finish_async().await?; + + Ok(file.map(|f| (f, max_record_batch_size))) } + .await; - let file = in_progress_file.finish()?; + if result.is_err() + && let Err(error) = in_progress_file.abort_async().await + { + debug!("Failed to abort in-progress stream spill: {error}"); + } - Ok(file.map(|f| (f, max_record_batch_size))) + result } /// Reads a spill file as a stream. The file must be created by the current @@ -255,14 +267,89 @@ mod tests { use crate::common::collect; use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; use crate::spill::{get_record_batch_memory_size, spill_manager::GetSlicedSize}; + use crate::stream::RecordBatchStreamAdapter; use arrow::datatypes::{DataType, Field, Schema}; use arrow::{ array::{ArrayRef, Int32Array, StringArray, StringViewArray}, record_batch::RecordBatch, }; - use datafusion_common::Result; - use datafusion_execution::runtime_env::RuntimeEnv; + use async_trait::async_trait; + use bytes::Bytes; + use datafusion_common::{DataFusionError, Result, not_impl_err}; + use datafusion_execution::disk_manager::DiskManagerBuilder; + use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + use datafusion_execution::{ + AsyncSpillWriter, SendableRecordBatchStream, SpillFile, SpillWriter, + TempFileFactory, + }; + use futures::Stream; + use std::path::Path; + use std::pin::Pin; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct AbortTrackingTempFileFactory { + abort_count: Arc, + } + + impl TempFileFactory for AbortTrackingTempFileFactory { + fn create_temp_file(&self, _description: &str) -> Result> { + Ok(Arc::new(AbortTrackingSpillFile { + abort_count: Arc::clone(&self.abort_count), + })) + } + } + + struct AbortTrackingSpillFile { + abort_count: Arc, + } + + #[async_trait] + impl SpillFile for AbortTrackingSpillFile { + fn path(&self) -> Option<&Path> { + None + } + + fn size(&self) -> Option { + Some(0) + } + + fn read_stream( + &self, + ) -> Result> + Send>>> { + Ok(Box::pin(futures::stream::empty())) + } + + fn open_writer(&self) -> Result> { + not_impl_err!("test backend only supports asynchronous writes") + } + + async fn open_async_writer(&self) -> Result> { + Ok(Box::new(AbortTrackingWriter { + abort_count: Arc::clone(&self.abort_count), + })) + } + } + + struct AbortTrackingWriter { + abort_count: Arc, + } + + #[async_trait] + impl AsyncSpillWriter for AbortTrackingWriter { + async fn write_all(&mut self, _data: Bytes) -> Result<()> { + Ok(()) + } + + async fn finish(&mut self) -> Result<()> { + Ok(()) + } + + async fn abort(&mut self) -> Result<()> { + self.abort_count.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + } fn build_test_spill_manager( env: Arc, @@ -357,6 +444,46 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_stream_spill_aborts_after_input_error() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let abort_count = Arc::new(AtomicUsize::new(0)); + let disk_manager_builder = DiskManagerBuilder::default().with_temp_file_factory( + Arc::new(AbortTrackingTempFileFactory { + abort_count: Arc::clone(&abort_count), + }), + ); + let env = RuntimeEnvBuilder::new() + .with_disk_manager_builder(disk_manager_builder) + .build_arc()?; + let manager = build_test_spill_manager(env, Arc::clone(&schema)); + let batch = build_writer_batch(Arc::clone(&schema))?; + let input = futures::stream::iter(vec![ + Ok(batch), + Err(DataFusionError::Execution( + "injected input stream failure".to_string(), + )), + ]); + let mut stream: SendableRecordBatchStream = + Box::pin(RecordBatchStreamAdapter::new(schema, input)); + + let error = manager + .spill_record_batch_stream_and_return_max_batch_memory( + &mut stream, + "abort-test", + ) + .await + .err() + .expect("the injected stream error should be returned"); + + assert!(error.to_string().contains("injected input stream failure")); + assert_eq!(abort_count.load(Ordering::Relaxed), 1); + Ok(()) + } + #[test] fn check_sliced_size_for_string_view_array() -> Result<()> { let array_length = 50; From 29bab9cb0af1874e0e47762db491fbce840fa7b8 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Sun, 6 Sep 2026 15:43:14 +0800 Subject: [PATCH 2/4] fix: avoid timers in spill drop cleanup Embedders may use Tokio runtimes without a time driver. Let backends bound abort latency while retaining the cleanup limit. Refs #23247 --- datafusion/execution/src/spill_file.rs | 4 ++++ datafusion/physical-plan/Cargo.toml | 2 +- datafusion/physical-plan/src/spill/mod.rs | 23 +++++++++-------------- 3 files changed, 14 insertions(+), 15 deletions(-) diff --git a/datafusion/execution/src/spill_file.rs b/datafusion/execution/src/spill_file.rs index d63d9e622e2eb..e73fbec08570e 100644 --- a/datafusion/execution/src/spill_file.rs +++ b/datafusion/execution/src/spill_file.rs @@ -97,6 +97,10 @@ pub trait AsyncSpillWriter: Send { /// Multipart backends should override this method to explicitly abort an /// in-progress upload. A successful call is terminal. DataFusion may retry /// an abort that returns an error, including during best-effort drop cleanup. + /// + /// Implementations should bound their own abort latency to avoid holding + /// DataFusion's limited cleanup slots indefinitely. DataFusion does not + /// impose a timeout because the caller's runtime may not have a time driver. async fn abort(&mut self) -> Result<()> { Ok(()) } diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 0aa22653b44ee..751b3935c5e96 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -88,7 +88,7 @@ num-traits = { workspace = true } parking_lot = { workspace = true } pin-project-lite = { workspace = true } serde_json = { workspace = true, features = ["preserve_order"] } -tokio = { workspace = true } +tokio = { workspace = true, features = ["time"] } [dev-dependencies] arrow-data = { workspace = true } diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 49197326e4917..3f56d2a6e0aaa 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -32,7 +32,6 @@ use std::collections::VecDeque; use std::pin::Pin; use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; -use std::time::Duration; use arrow::array::{ Array, ArrayRef, BinaryViewArray, BufferSpec, GenericByteViewArray, StringViewArray, @@ -520,11 +519,10 @@ struct AsyncIPCStreamWriter { abort_handle: Option, } -/// Bounds detached cleanup work after query cancellation. If all permits are -/// occupied, the writer is dropped and the backend lifecycle policy is the -/// remaining cleanup mechanism. +/// Bounds the number of detached cleanup tasks after query cancellation. If all +/// permits are occupied, the writer is dropped and the backend lifecycle policy +/// is the remaining cleanup mechanism. const MAX_CONCURRENT_SPILL_ABORTS: usize = 8; -const SPILL_ABORT_TIMEOUT: Duration = Duration::from_secs(30); static SPILL_ABORT_PERMITS: LazyLock> = LazyLock::new(|| Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_SPILL_ABORTS))); @@ -649,16 +647,12 @@ impl Drop for AsyncIPCStreamWriter { }; let _abort_task = handle.spawn(async move { - match tokio::time::timeout(SPILL_ABORT_TIMEOUT, writer.abort()).await { - Ok(Ok(())) => {} - Ok(Err(error)) => { - debug!("Failed to abort dropped spill writer: {error}"); - } - Err(_) => { - debug!("Timed out aborting dropped spill writer"); - } + // Backends bound their own abort latency; the runtime may not have + // a time driver. See the AsyncSpillWriter::abort contract. + if let Err(error) = writer.abort().await { + debug!("Failed to abort dropped spill writer: {error}"); } - // `permit` is released when this bounded cleanup task exits. + // Keep the permit until cleanup finishes. drop(permit); }); } @@ -936,6 +930,7 @@ mod tests { use std::io::Cursor; use std::sync::Mutex; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + use std::time::Duration; use tokio::sync::Notify; struct YieldingAsyncWriter { From 7b7a39634393ea1969a887c54cce2c244fd0529f Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 7 Sep 2026 11:16:52 +0800 Subject: [PATCH 3/4] fix: borrow decimal values in SQL tests Decimal formatting only borrows its input. Preserve that ownership contract to satisfy all-feature Clippy checks. --- datafusion/sqllogictest/src/engines/conversion.rs | 2 +- datafusion/sqllogictest/src/engines/postgres_engine/mod.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/sqllogictest/src/engines/conversion.rs b/datafusion/sqllogictest/src/engines/conversion.rs index d22b518234803..13178e1db2c1e 100644 --- a/datafusion/sqllogictest/src/engines/conversion.rs +++ b/datafusion/sqllogictest/src/engines/conversion.rs @@ -96,7 +96,7 @@ pub(crate) fn arrow_decimal_to_str( } #[cfg(feature = "postgres")] -pub(crate) fn decimal_to_str(value: BigDecimal) -> String { +pub(crate) fn decimal_to_str(value: &BigDecimal) -> String { value.to_plain_string() } diff --git a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs index 7ab5f1977d0e4..daa5acf165891 100644 --- a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs +++ b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs @@ -377,7 +377,7 @@ fn cell_to_string(row: &SimpleQueryRow, column_type: &Type, idx: usize) -> Strin (&Type::INT4, Some(value)) => value.parse::().unwrap().to_string(), (&Type::INT8, Some(value)) => value.parse::().unwrap().to_string(), (&Type::NUMERIC, Some(value)) => { - decimal_to_str(BigDecimal::from_str(value).unwrap()) + decimal_to_str(&BigDecimal::from_str(value).unwrap()) } // Parse date/time strings explicitly to avoid locale-specific formatting. (&Type::DATE, Some(value)) => NaiveDate::parse_from_str(value, "%Y-%m-%d") From 33f4ab2c8e88e5cf82870970a748f747b5eb9ee8 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 7 Sep 2026 11:17:05 +0800 Subject: [PATCH 4/4] fix: respect async spill memory budgets Async spill writes must retain their input budget without bypassing pool limits. Reuse reserved workspace and release consumed merge inputs so spilling can progress within its budget. Refs #23247 --- datafusion/physical-plan/src/sorts/builder.rs | 30 +++- datafusion/physical-plan/src/sorts/sort.rs | 158 ++++++++++-------- .../src/sorts/sort/spill_tests.rs | 6 +- 3 files changed, 120 insertions(+), 74 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..77d5306c01f62 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -168,12 +168,15 @@ impl BatchBuilder { // is finished. This means all remaining rows from all but the last batch // for each stream have been yielded to the newly created record batch // - // We can therefore drop all but the last batch for each stream + // Keep the last batch only if its cursor still has unread rows. + // Fully-consumed batches must release their reservation before the + // output consumer reserves memory for the newly built batch. let mut batch_idx = 0; let mut retained = 0; self.batches.retain(|(stream_idx, batch)| { let stream_cursor = &mut self.cursors[*stream_idx]; - let retain = stream_cursor.batch_idx == batch_idx; + let retain = stream_cursor.batch_idx == batch_idx + && stream_cursor.row_idx < batch.num_rows(); batch_idx += 1; if retain { @@ -341,6 +344,29 @@ mod tests { assert!(matches!(error, DataFusionError::Execution(msg) if msg == "boom")); } + #[test] + fn test_releases_fully_consumed_last_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("test").register(&pool); + let mut builder = BatchBuilder::new(Arc::clone(&schema), 1, 1, reservation); + for values in [[1, 2], [3, 4]] { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values.to_vec()))], + )?; + let size = get_record_batch_memory_size(&batch); + builder.push_batch(0, batch)?; + builder.push_row(0); + assert_eq!(builder.build_record_batch()?.unwrap().num_rows(), 1); + assert_eq!(pool.reserved(), size, "unread rows must remain reserved"); + builder.push_row(0); + assert_eq!(builder.build_record_batch()?.unwrap().num_rows(), 1); + assert_eq!(pool.reserved(), 0, "consumed batches must release memory"); + } + Ok(()) + } + #[test] fn test_try_interleave_columns_surfaces_arrow_offset_overflow() { let batch = overflow_list_batch(); diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 1cbd391fd7534..6f7b8aa42a7ec 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -419,16 +419,8 @@ impl ExternalSorter { self.metrics.spill_metrics.spill_file_count.value() } - /// Appending globally sorted batches to the in-progress spill file, and clears - /// the `globally_sorted_batches` (also its memory reservation) afterwards. - async fn consume_and_spill_append( - &mut self, - globally_sorted_batches: &mut Vec, - ) -> Result<()> { - if globally_sorted_batches.is_empty() { - return Ok(()); - } - + /// Appends a globally sorted batch, retaining its reservation until written. + async fn consume_and_spill_append(&mut self, batch: RecordBatch) -> Result<()> { // Lazily initialize the in-progress spill file if self.in_progress_spill_file.is_none() { self.in_progress_spill_file = @@ -437,8 +429,7 @@ impl ExternalSorter { debug!("Spilling sort data of ExternalSorter to disk whilst inserting"); - let batches_to_spill = std::mem::take(globally_sorted_batches); - // Keep the reservation alive while the batches remain in memory across + // Keep the reservation alive while the batch remains in memory across // asynchronous writes. It is released on success or error via RAII. let _spill_reservation = self.reservation.take(); @@ -447,16 +438,8 @@ impl ExternalSorter { internal_datafusion_err!("In-progress spill file should be initialized") })?; - for batch in batches_to_spill { - let gc_sliced_size = in_progress_file.append_batch_async(&batch).await?; - - *max_record_batch_size = (*max_record_batch_size).max(gc_sliced_size); - } - - assert_or_internal_err!( - globally_sorted_batches.is_empty(), - "This function consumes globally_sorted_batches, so it should be empty after taking." - ); + let gc_sliced_size = in_progress_file.append_batch_async(&batch).await?; + *max_record_batch_size = (*max_record_batch_size).max(gc_sliced_size); Ok(()) } @@ -520,49 +503,35 @@ impl ExternalSorter { self.in_mem_batches.is_empty(), "in_mem_batches should be empty after constructing sorted stream" ); - // 'global' here refers to all buffered batches when the memory limit is - // reached. This variable will buffer the sorted batches after - // sort-preserving merge and incrementally append to spill files. - let mut globally_sorted_batches: Vec = vec![]; - while let Some(batch) = sorted_stream.next().await { let batch = batch?; - let sorted_size = get_reserved_bytes_for_record_batch(&batch)?; - let reservation_failed = match self.reservation.try_grow(sorted_size) { - Ok(()) => false, + // Sorting is complete: retain the batch's footprint, not the input + // estimate that also budgets for creating a sorted copy. + let sorted_size = get_record_batch_memory_size(&batch); + let spill_workspace = match self.reservation.try_grow(sorted_size) { + Ok(()) => None, Err(_) => { - // The batch is already materialized, so account for it while - // spilling even if this temporarily exceeds the pool limit. - self.reservation.grow(sorted_size); - true + // Reuse already-reserved workspace without bypassing the + // execution pool's limit. Any remainder still needs a grant + // through the original sort consumer. + let workspace = self.merge_pool.borrow(sorted_size); + self.reservation.try_grow(sorted_size - workspace.size())?; + Some(workspace) } }; - // Even if the reservation is not enough, the batch is already in - // memory, so it's okay to combine it with previously sorted - // batches, and spill together. - globally_sorted_batches.push(batch); - if reservation_failed { - self.consume_and_spill_append(&mut globally_sorted_batches) - .await?; // reservation is released when the spill completes - } + // Write each batch before polling the merge again. Accumulating + // output would compete with the merge's workspace without combining + // any writes. Keep both forms of reservation alive across the await. + self.consume_and_spill_append(batch).await?; + drop(spill_workspace); } // Drop early to free up memory reserved by the sorted stream, otherwise the // upcoming `self.reserve_memory_for_merge()` may fail due to insufficient memory. drop(sorted_stream); - self.consume_and_spill_append(&mut globally_sorted_batches) - .await?; self.spill_finish().await?; - // Sanity check after spilling - let buffers_cleared_property = - self.in_mem_batches.is_empty() && globally_sorted_batches.is_empty(); - assert_or_internal_err!( - buffers_cleared_property, - "in_mem_batches and globally_sorted_batches should be cleared before" - ); - // Reserve headroom for next sort/merge self.reserve_memory_for_merge()?; @@ -3386,6 +3355,32 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_spill_output_respects_memory_limit() -> Result<()> { + let result = test_sort_output_batch_size_and_base_metrics(10, 25, |batches| { + let batches_memory = batches.iter().map(|b| b.get_array_memory_size()).sum(); + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(100) + .with_sort_in_place_threshold_bytes(1) + .with_sort_spill_reservation_bytes(1), + ) + .with_runtime( + RuntimeEnvBuilder::default() + .with_memory_limit(batches_memory, 1.0) + .build_arc() + .unwrap(), + ) + }) + .await; + assert!(matches!( + result, + Err(DataFusionError::ResourcesExhausted(_)) + )); + Ok(()) + } + #[tokio::test] async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_to_spill() -> Result<()> { @@ -3396,6 +3391,11 @@ mod tests { .iter() .map(|b| b.get_array_memory_size()) .sum::(); + // Leave space for one output batch while the merge still holds + // partially consumed input batches. The insufficient-budget case + // is covered by test_spill_output_respects_memory_limit. + let spill_workspace = + make_partition(batch_size as i32).get_array_memory_size(); TaskContext::default() .with_session_config( @@ -3403,11 +3403,11 @@ mod tests { .with_batch_size(batch_size) // To make sure there is no in place sorting .with_sort_in_place_threshold_bytes(1) - .with_sort_spill_reservation_bytes(1), + .with_sort_spill_reservation_bytes(spill_workspace), ) .with_runtime( RuntimeEnvBuilder::default() - .with_memory_limit(batches_memory, 1.0) + .with_memory_limit(batches_memory + spill_workspace, 1.0) .build_arc() .unwrap(), ) @@ -3957,7 +3957,27 @@ mod tests { #[tokio::test] async fn test_spill_reservation_held_during_async_write() -> Result<()> { - let pool: Arc = Arc::new(GreedyMemoryPool::new(0)); + let schema = Arc::new(Schema::new(vec![Field::new( + "x", + DataType::Utf8View, + false, + )])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringViewArray::from_iter_values( + (0..4096).rev().map(|i| format!("{i:08}{}", "x".repeat(87))), + ))], + )?; + let ordering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(); + let batch_size = 1024; + let sorted_bytes = sort_batch_chunked(&batch, &ordering, batch_size)? + .iter() + .map(get_record_batch_memory_size) + .sum::(); + let workspace_bytes = sorted_bytes - get_reserved_bytes_for_record_batch(&batch)?; + assert!(workspace_bytes > 0); + let pool = spill_tests::AdjustablePool::new(sorted_bytes); let write_started = Arc::new(Notify::new()); let aborted = Arc::new(Notify::new()); let abort_count = Arc::new(AtomicUsize::new(0)); @@ -3969,29 +3989,26 @@ mod tests { }), ); let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&pool)) + .with_memory_pool(Arc::clone(&pool) as Arc) .with_disk_manager_builder(disk_manager_builder) .build_arc()?; - let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); let metrics = ExecutionPlanMetricsSet::new(); let mut sorter = ExternalSorter::new( 0, - Arc::clone(&schema), - [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(), - 128, + schema, + ordering, + batch_size, + workspace_bytes, 0, - usize::MAX, SpillCompression::Uncompressed, &metrics, runtime, )?; - let batch = RecordBatch::try_new( - schema, - vec![Arc::new(Int32Array::from(vec![3, 2, 1]))], - )?; - let reserved_bytes = get_reserved_bytes_for_record_batch(&batch)?; - sorter.reservation.grow(reserved_bytes); - sorter.in_mem_batches.push(batch); + sorter.insert_batch(batch).await?; + // Existing reservations remain valid, but the emitted batch must reuse + // workspace rather than request new parent capacity. + pool.set_limit(sorted_bytes - 1); + let merge_pool = Arc::clone(&sorter.merge_pool); #[expect(clippy::disallowed_methods)] // spawn allowed only in tests let task = @@ -4008,9 +4025,12 @@ mod tests { let _ = task.await; panic!("spill write did not start before the timeout"); } + // Idle workspace must be releasable without releasing the loan held by + // the pending write. The stream still owns the remaining sorted batches. + merge_pool.release_unused(); assert_eq!( pool.reserved(), - reserved_bytes, + sorted_bytes, "resident batches must remain accounted for while an async spill is pending" ); diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index f41ca4507c25a..3c97b5ad65f05 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -57,13 +57,13 @@ struct AllocationState { /// Existing allocations survive a lower limit, but fresh allocations must fit. /// The test changes the limit before insertion, independently of spill internals. #[derive(Debug)] -struct AdjustablePool { +pub(super) struct AdjustablePool { capacity: usize, state: Mutex, } impl AdjustablePool { - fn new(capacity: usize) -> Arc { + pub(super) fn new(capacity: usize) -> Arc { Arc::new(Self { capacity, state: Mutex::new(AllocationState { @@ -73,7 +73,7 @@ impl AdjustablePool { }) } - fn set_limit(&self, limit: usize) { + pub(super) fn set_limit(&self, limit: usize) { assert!(limit <= self.capacity); self.state.lock().unwrap().limit = limit; }