This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23353-4f3980210c909aa1e91fa5b92f6899f78bec32be in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit f269b94bc2411e55240322417791c78446fd00a7 Author: pantShrey <[email protected]> AuthorDate: Sat Jul 11 05:55:52 2026 +0530 test: add Poll::Pending spill stream coverage for async spill re-entry paths (#23353) ## Which issue does this PR close? - Closes #22879. ## Rationale for this change The SortMergeJoin refactor in #22230 introduced async spill read paths in both `materializing_stream.rs` (`poll_spilled_batches`) and `bitwise_stream.rs` (the spilled-batch loop in `process_key_match_with_filter`). Both call `SpillFile::read_stream()` and poll the resulting byte stream, which can return `Poll::Pending`. The existing spill tests exercise spilling and restoring batches, but always run against the default local-file `SpillFile`, whose `read_stream()` happens to resolve synchronously on first poll in practice. As a result, the re-entry logic in both streams has no dedicated test coverage. ## What changes are included in this PR? - Adds a `PendingSpillFile` / `PendingTempFileFactory` pair in `sort_merge_join/tests.rs` that wraps the default local `SpillFile` backend (via `DiskManagerMode::Custom`). Every spill read splits the real spill file's bytes into small fixed-size chunks and yields `Poll::Pending` before each chunk, forcing `SpillReaderStream`'s IPC decoder to genuinely suspend mid-read and resume from partially buffered state across multiple real Pending cycles, rather than completing on the first poll. - Adds `materializing_spill_pending_stream`, covering Inner/Left/ Right/Full joins (the `BufferedBatchState::Spilled` restore path in `poll_spilled_batches`). - Adds `bitwise_spill_pending_stream`, covering LeftSemi/LeftAnti/ RightSemi/RightAnti joins with a filter (the `inner_key_spill` re-entry path in `process_key_match_with_filter`). - Both tests compare the Pending-forced spilled result against a no-spill run of the same join, asserting the results are identical, so any state corruption introduced by suspending mid-read would be caught. ## Are these changes tested? Yes -- this is a PR to test recent changes ## Are there any user-facing changes? No. --- .../src/joins/sort_merge_join/tests.rs | 306 ++++++++++++++++++++- 1 file changed, 305 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index bc8b15472a..ccd9c155f1 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -49,6 +49,7 @@ use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch}; use arrow::datatypes::{DataType, Field, Schema}; use arrow_ord::sort::SortColumn; use arrow_schema::SchemaRef; +use bytes::Bytes; use datafusion_common::JoinType::*; use datafusion_common::{ JoinSide, internal_err, @@ -59,9 +60,12 @@ use datafusion_common::{ }; use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; -use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::disk_manager::{ + DiskManager, DiskManagerBuilder, DiskManagerMode, +}; use datafusion_execution::memory_pool::MemoryConsumer; use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_execution::spill_file::{SpillFile, SpillWriter, TempFileFactory}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::BinaryExpr; @@ -70,6 +74,7 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use futures::{Stream, StreamExt}; use insta::assert_snapshot; use itertools::Itertools; +use std::collections::VecDeque; fn build_table( a: (&str, &Vec<i32>), @@ -5065,3 +5070,302 @@ async fn spill_read_back_single_source() -> Result<()> { Ok(()) } + +/// Small chunk size so even tiny test spill files are split into several +/// pieces, forcing multiple genuine suspend/resume cycles instead of one. +const PENDING_CHUNK_SIZE: usize = 16; + +/// Splits real spill bytes into fixed-size chunks and yields `Poll::Pending` +/// before every chunk +struct PendingChunkedStream { + chunks: VecDeque<Bytes>, + yield_pending: bool, +} + +impl PendingChunkedStream { + fn new(bytes: Bytes) -> Self { + let mut chunks = VecDeque::new(); + if bytes.is_empty() { + chunks.push_back(bytes); + } else { + let mut remaining = bytes; + while !remaining.is_empty() { + let take = PENDING_CHUNK_SIZE.min(remaining.len()); + chunks.push_back(remaining.split_to(take)); + } + } + Self { + chunks, + yield_pending: true, + } + } +} + +impl Stream for PendingChunkedStream { + type Item = Result<Bytes>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.yield_pending { + self.yield_pending = false; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + // Pending before every subsequent chunk as well. + self.yield_pending = true; + match self.chunks.pop_front() { + Some(chunk) => Poll::Ready(Some(Ok(chunk))), + None => Poll::Ready(None), + } + } +} + +/// A `SpillFile` that delegates everything to a real local spill file, +/// except `read_stream`, which is forced through `PendingChunkedStream`. +struct PendingSpillFile { + inner: Arc<dyn SpillFile>, +} + +impl SpillFile for PendingSpillFile { + fn path(&self) -> Option<&std::path::Path> { + self.inner.path() + } + + fn size(&self) -> Option<u64> { + self.inner.size() + } + + fn read_stream(&self) -> Result<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>> { + let path = self + .inner + .path() + .expect("PendingSpillFile only wraps local files") + .to_owned(); + + let stream = futures::stream::once(async move { + tokio::fs::read(&path) + .await + .map(Bytes::from) + .map_err(datafusion_common::DataFusionError::IoError) + }) + .flat_map( + |read_result| -> Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>> { + match read_result { + Ok(bytes) => Box::pin(PendingChunkedStream::new(bytes)), + Err(e) => Box::pin(futures::stream::once(async move { Err(e) })), + } + }, + ); + + Ok(Box::pin(stream)) + } + + fn open_writer(&self) -> Result<Box<dyn SpillWriter>> { + self.inner.open_writer() + } +} + +/// Wraps the default `OsTmpDirectory` factory so every spill file it +/// creates is a [`PendingSpillFile`]. +struct PendingTempFileFactory { + inner: Arc<DiskManager>, +} + +impl TempFileFactory for PendingTempFileFactory { + fn create_temp_file(&self, description: &str) -> Result<Arc<dyn SpillFile>> { + Ok(Arc::new(PendingSpillFile { + inner: self.inner.create_tmp_file(description)?, + })) + } +} + +fn pending_disk_manager_builder() -> DiskManagerBuilder { + let inner = Arc::new( + DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .build() + .unwrap(), + ); + DiskManagerBuilder::default().with_mode(DiskManagerMode::Custom(Arc::new( + PendingTempFileFactory { inner }, + ))) +} + +/// Materializing-side (Inner/Left/Right/Full) coverage: identical to +/// `overallocation_multi_batch_spill`, but every spill read goes through +/// `PendingSpillFile`, so `poll_spilled_batches` must actually hit and +/// recover from `Poll::Pending` mid-read. +#[tokio::test] +async fn materializing_spill_pending_stream() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![1, 1]), + ("c1", &vec![6, 7]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10]), + ("b2", &vec![1, 1]), + ("c2", &vec![50, 60]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![20, 30]), + ("b2", &vec![1, 1]), + ("c2", &vec![70, 80]), + ); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2]); + let right = build_table_from_batches(vec![right_batch_1, right_batch_2]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(500, 1.0) + .with_disk_manager_builder(pending_disk_manager_builder()) + .build_arc()?; + + for join_type in [Inner, Left, Right, Full] { + let task_ctx = + Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "expected spill_count > 0 for {join_type:?}" + ); + + // Compare against a no-spill run to make sure the Pending + // re-entry path didn't corrupt or drop any data. + let task_ctx_no_spill = Arc::new(TaskContext::default()); + let join_no_spill = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + assert_eq!( + spilled_result, no_spill_result, + "Pending-forced spill read produced different results for {join_type:?}" + ); + } + + Ok(()) +} + +/// Bitwise-side (Semi/Anti) coverage: identical to `bitwise_spill_with_filter`, +/// but every spill read goes through `PendingSpillFile`, forcing +/// `process_key_match_with_filter`'s spilled-batch loop to actually hit and +/// resume from `Poll::Pending`. +#[tokio::test] +async fn bitwise_spill_pending_stream() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4, 5, 6]), + ("b1", &vec![1, 2, 3, 4, 5, 6]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40, 50]), + ("b1", &vec![1, 3, 4, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + // c1 < c2 is always true for matching keys — same filter as + // bitwise_spill_with_filter, so the inner key group is buffered + // (and spilled) rather than short-circuited. + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Lt, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder(pending_disk_manager_builder()) + .build_arc()?; + + for join_type in [LeftSemi, LeftAnti, RightSemi, RightAnti] { + let task_ctx = + Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let join = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "expected spill_count > 0 for {join_type:?}" + ); + + let task_ctx_no_spill = Arc::new(TaskContext::default()); + let join_no_spill = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + assert_eq!( + spilled_result, no_spill_result, + "Pending-forced spill read produced different results for {join_type:?}" + ); + } + + Ok(()) +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
