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..e73fbec08570e 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,93 @@ 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. + /// + /// 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(()) + } +} + +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/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/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 490ea7cc85776..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. - 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,35 +429,28 @@ 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 batch remains 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(|| { 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(&batch)?; - - *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(()) } /// 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 +462,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" @@ -501,38 +503,34 @@ 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 = self.reservation.try_grow(sorted_size).is_err(); - // 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() - } + // 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(_) => { + // 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) + } + }; + // 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)?; - self.spill_finish()?; - - // 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" - ); + self.spill_finish().await?; // Reserve headroom for next sort/merge self.reserve_memory_for_merge()?; @@ -1490,8 +1488,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 +2122,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 +2142,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 { @@ -3262,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<()> { @@ -3272,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( @@ -3279,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(), ) @@ -3831,6 +3955,95 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_spill_reservation_held_during_async_write() -> Result<()> { + 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)); + 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) as Arc) + .with_disk_manager_builder(disk_manager_builder) + .build_arc()?; + let metrics = ExecutionPlanMetricsSet::new(); + let mut sorter = ExternalSorter::new( + 0, + schema, + ordering, + batch_size, + workspace_bytes, + 0, + SpillCompression::Uncompressed, + &metrics, + runtime, + )?; + 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 = + 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"); + } + // 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(), + sorted_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/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; } 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..3f56d2a6e0aaa 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,7 +30,7 @@ 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 arrow::array::{ @@ -43,14 +43,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 +466,205 @@ 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 finish(mut self) -> Result<()> { - self.inner.finish() + pub fn flush(&mut self) -> Result<()> { + use std::io::Write; + self.writer.flush()?; + Ok(()) + } + + 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, +} + +/// 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; +static SPILL_ABORT_PERMITS: LazyLock> = + LazyLock::new(|| Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_SPILL_ABORTS))); + +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(), + } + } + + 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) + } - self.total_bytes_written += n; + pub async fn abort(&mut self) -> Result<()> { + let Some(writer) = self.writer.as_mut() else { + return Ok(()); + }; + writer.abort().await?; + self.writer = None; + Ok(()) + } - Ok(n) + async fn abort_after_error(&mut self) { + if let Err(error) = self.abort().await { + debug!("Failed to abort spill writer: {error}"); + } } +} - fn flush(&mut self) -> std::io::Result<()> { - self.inner.flush() +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 { + // 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}"); + } + // Keep the permit until cleanup finishes. + drop(permit); + }); } } -/// 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, +/// Encodes Arrow IPC streams without performing I/O. +struct IPCStreamEncoder { + encoder: Option, } -impl IPCStreamWriter { - /// Create new writer +impl IPCStreamEncoder { + /// Create a new encoder. /// /// # Codec contract /// @@ -528,11 +676,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 +690,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 +924,309 @@ 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 std::time::Duration; + 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 +1858,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; 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")