From a12fda0717e30065b4ce83f170fd5e1cb0b6d92a Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Sun, 9 Aug 2026 13:32:00 -0400 Subject: [PATCH 01/13] refactor: CrossJoinStream using generators --- .../physical-plan/src/joins/cross_join.rs | 23 ++----------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 16155aaafdd9c..730700a069d74 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -570,8 +570,8 @@ struct CrossJoinStream { left_index: usize, /// Join execution metrics join_metrics: BuildProbeJoinMetrics, - /// State of the stream - state: CrossJoinStreamState, + /// Currently processed right side batch + state: Option, /// Left data (copy of the entire buffered left side) left_data: RecordBatch, /// Batch transformer @@ -584,25 +584,6 @@ impl RecordBatchStream for CrossJoinStream Result<&RecordBatch> { - match self { - CrossJoinStreamState::BuildBatches(rb) => Ok(rb), - _ => internal_err!("Expected RecordBatch in BuildBatches state"), - } - } -} - fn build_batch( left_index: usize, batch: &RecordBatch, From 1d7173218ebeba0c38d02ad4cac800b5141adfe5 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Sun, 9 Aug 2026 13:40:08 -0400 Subject: [PATCH 02/13] removing polling methods --- .../physical-plan/src/joins/cross_join.rs | 20 +------------------ 1 file changed, 1 insertion(+), 19 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 730700a069d74..b179596502007 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -571,19 +571,13 @@ struct CrossJoinStream { /// Join execution metrics join_metrics: BuildProbeJoinMetrics, /// Currently processed right side batch - state: Option, + processed_right_data: Option, /// Left data (copy of the entire buffered left side) left_data: RecordBatch, /// Batch transformer batch_transformer: T, } -impl RecordBatchStream for CrossJoinStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} - fn build_batch( left_index: usize, batch: &RecordBatch, @@ -612,18 +606,6 @@ fn build_batch( .map_err(Into::into) } -#[async_trait] -impl Stream for CrossJoinStream { - type Item = Result; - - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> Poll> { - self.poll_next_impl(cx) - } -} - impl CrossJoinStream { /// Separate implementation function that unpins the [`CrossJoinStream`] so /// that partial borrows work correctly From bd1bd44486d75d3c398b5512644fc4b99810016d Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Mon, 10 Aug 2026 18:52:05 -0400 Subject: [PATCH 03/13] moved fetch_probe_batch to be async --- .../physical-plan/src/joins/cross_join.rs | 28 +++++++++---------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index b179596502007..df3186fb9aa29 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -358,7 +358,7 @@ impl ExecutionPlan for CrossJoinExec { right: stream, left_index: 0, join_metrics, - state: CrossJoinStreamState::WaitBuildSide, + processed_right_data: None, left_data: RecordBatch::new_empty(self.left().schema()), batch_transformer: BatchSplitter::new(batch_size), })) @@ -369,7 +369,7 @@ impl ExecutionPlan for CrossJoinExec { right: stream, left_index: 0, join_metrics, - state: CrossJoinStreamState::WaitBuildSide, + processed_right_data: None, left_data: RecordBatch::new_empty(self.left().schema()), batch_transformer: NoopBatchTransformer::new(), })) @@ -609,10 +609,10 @@ fn build_batch( impl CrossJoinStream { /// Separate implementation function that unpins the [`CrossJoinStream`] so /// that partial borrows work correctly - fn poll_next_impl( + async fn next( &mut self, cx: &mut std::task::Context<'_>, - ) -> Poll>> { + ) -> Option> { loop { return match self.state { CrossJoinStreamState::WaitBuildSide => { @@ -634,7 +634,7 @@ impl CrossJoinStream { fn collect_build_side( &mut self, cx: &mut std::task::Context<'_>, - ) -> Poll>>> { + ) -> Result>> { let build_timer = self.join_metrics.build_time.timer(); let left_data = match ready!(self.left_fut.get(cx)) { Ok(left_data) => left_data, @@ -647,34 +647,32 @@ impl CrossJoinStream { StatefulStreamResult::Ready(None) } else { self.left_data = left_data; - self.state = CrossJoinStreamState::FetchProbeBatch; StatefulStreamResult::Continue }; - Poll::Ready(Ok(result)) + Ok(result) } /// Fetches the probe (right) batch, updates the metrics, and save the batch in the state. /// Then, the state is updated to build result batches. - fn fetch_probe_batch( + async fn fetch_probe_batch( &mut self, - cx: &mut std::task::Context<'_>, - ) -> Poll>>> { + ) -> Result>> { self.left_index = 0; - let right_data = match ready!(self.right.poll_next_unpin(cx)) { + let right_data = match self.right.next().await { Some(Ok(right_data)) => right_data, - Some(Err(e)) => return Poll::Ready(Err(e)), + Some(Err(e)) => return Err(e), None => { // Release the right (probe) input pipeline's resources. let right_schema = self.right.schema(); self.right = Box::pin(EmptyRecordBatchStream::new(right_schema)); - return Poll::Ready(Ok(StatefulStreamResult::Ready(None))); + return Ok(StatefulStreamResult::Ready(None)); } }; self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(right_data.num_rows()); - self.state = CrossJoinStreamState::BuildBatches(right_data); - Poll::Ready(Ok(StatefulStreamResult::Continue)) + self.processed_right_data = Some(right_data); + Ok(StatefulStreamResult::Continue) } /// Joins the indexed row of left data with the current probe batch. From 8ca90e42b0f175afa7f832205d745377aa3d6901 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Mon, 10 Aug 2026 19:20:44 -0400 Subject: [PATCH 04/13] comments changed --- .../physical-plan/src/joins/cross_join.rs | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index df3186fb9aa29..54ee28a827833 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -44,7 +44,8 @@ use arrow::compute::concat_batches; use arrow::datatypes::{Fields, Schema, SchemaRef}; use datafusion_common::stats::Precision; use datafusion_common::{ - JoinType, Result, ScalarValue, assert_eq_or_internal_err, internal_err, + DataFusionError, JoinType, Result, ScalarValue, assert_eq_or_internal_err, + internal_err, }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; @@ -631,23 +632,21 @@ impl CrossJoinStream { /// Collects build (left) side of the join into the state. In case of an empty build batch, /// the execution terminates. Otherwise, the state is updated to fetch probe (right) batch. - fn collect_build_side( - &mut self, - cx: &mut std::task::Context<'_>, - ) -> Result>> { + /// Returns true if build side was loaded and non-empty + fn collect_build_side(&mut self, cx: &mut std::task::Context<'_>) -> Result { let build_timer = self.join_metrics.build_time.timer(); let left_data = match ready!(self.left_fut.get(cx)) { Ok(left_data) => left_data, - Err(e) => return Poll::Ready(Err(e)), + Err(e) => return Err(e), }; build_timer.done(); let left_data = left_data.merged_batch.clone(); let result = if left_data.num_rows() == 0 { - StatefulStreamResult::Ready(None) + false } else { self.left_data = left_data; - StatefulStreamResult::Continue + true }; Ok(result) } @@ -677,15 +676,19 @@ impl CrossJoinStream { /// Joins the indexed row of left data with the current probe batch. /// If all the results are produced, the state is set to fetch new probe batch. - fn build_batches(&mut self) -> Result>> { - let right_batch = self.state.try_as_record_batch()?; + async fn build_batches( + &mut self, + ) -> Result>> { + let right_batch = self + .processed_right_data + .ok_or(internal_err!("Expected RecordBatch for right side"))?; if self.left_index < self.left_data.num_rows() { match self.batch_transformer.next() { None => { let join_timer = self.join_metrics.join_time.timer(); let result = build_batch( self.left_index, - right_batch, + &right_batch, &self.left_data, &self.schema, ); From 2bf600260d89de831d425c65278ea914e13be109 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Tue, 11 Aug 2026 12:35:33 -0400 Subject: [PATCH 05/13] moving to single join method instead of polls --- .../physical-plan/src/joins/cross_join.rs | 71 ++++++++++--------- 1 file changed, 39 insertions(+), 32 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 54ee28a827833..e7789d84222b7 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -32,7 +32,7 @@ use crate::projection::{ physical_to_column_exprs, }; use crate::statistics::{ChildStats, StatisticsArgs}; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, @@ -47,8 +47,8 @@ use datafusion_common::{ DataFusionError, JoinType, Result, ScalarValue, assert_eq_or_internal_err, internal_err, }; -use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, async_try_stream}; use datafusion_physical_expr::equivalence::join_equivalence_properties; use async_trait::async_trait; @@ -352,29 +352,34 @@ impl ExecutionPlan for CrossJoinExec { )) })?; - if enforce_batch_size_in_joins { - Ok(Box::pin(CrossJoinStream { - schema: Arc::clone(&self.schema), - left_fut, - right: stream, - left_index: 0, - join_metrics, - processed_right_data: None, - left_data: RecordBatch::new_empty(self.left().schema()), - batch_transformer: BatchSplitter::new(batch_size), - })) - } else { - Ok(Box::pin(CrossJoinStream { - schema: Arc::clone(&self.schema), - left_fut, - right: stream, - left_index: 0, - join_metrics, - processed_right_data: None, - left_data: RecordBatch::new_empty(self.left().schema()), - batch_transformer: NoopBatchTransformer::new(), - })) - } + let mut state = CrossJoinStream { + schema: Arc::clone(&self.schema), + left_fut, + right: stream, + left_index: 0, + join_metrics, + processed_right_data: None, + left_data: RecordBatch::new_empty(self.left().schema()), + batch_transformer: if enforce_batch_size_in_joins { + BatchSplitter::new(batch_size) + } else { + NoopBatchTransformer::new() //todo: figure this out lol + }, + }; + + let stream = async_try_stream(|mut emitter| async move { + state.start_join_time(); + let result = state.join(&mut emitter).await; + state.stop_join_time(); + result + }); + // ObservedStream records the baseline metrics (output rows/batches, + // end time) + Ok(Box::pin(ObservedStream::new( + Box::pin(RecordBatchStreamAdapter::new(state.schema, stream)), + baseline_metrics, //todo: fix this too + None, + ))) } fn child_stats_requests(&self, partition: Option) -> Vec { @@ -608,12 +613,13 @@ fn build_batch( } impl CrossJoinStream { - /// Separate implementation function that unpins the [`CrossJoinStream`] so - /// that partial borrows work correctly - async fn next( - &mut self, - cx: &mut std::task::Context<'_>, - ) -> Option> { + async fn join(&mut self, cx: &mut std::task::Context<'_>) -> Result<()> { + if !self.collect_build_side(cx)? { + return Ok(()); + } + + while self.processed_right_data.is_none() {} // <- maybe add helper functions for when we have more data to determine the looping condition? + loop { return match self.state { CrossJoinStreamState::WaitBuildSide => { @@ -635,7 +641,8 @@ impl CrossJoinStream { /// Returns true if build side was loaded and non-empty fn collect_build_side(&mut self, cx: &mut std::task::Context<'_>) -> Result { let build_timer = self.join_metrics.build_time.timer(); - let left_data = match ready!(self.left_fut.get(cx)) { + let left_data = match self.left_fut.get(cx) { + //todo: figure out how to make this async maybe Ok(left_data) => left_data, Err(e) => return Err(e), }; From 503a69a88a2867ff0b20ac796be4528ac36b5b89 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Tue, 11 Aug 2026 18:40:25 -0400 Subject: [PATCH 06/13] removing context & batch transformers --- .../physical-plan/src/joins/cross_join.rs | 35 ++++++++++--------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index e7789d84222b7..30024f719852e 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -48,7 +48,7 @@ use datafusion_common::{ internal_err, }; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_execution::{TaskContext, async_try_stream}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use datafusion_physical_expr::equivalence::join_equivalence_properties; use async_trait::async_trait; @@ -360,16 +360,12 @@ impl ExecutionPlan for CrossJoinExec { join_metrics, processed_right_data: None, left_data: RecordBatch::new_empty(self.left().schema()), - batch_transformer: if enforce_batch_size_in_joins { - BatchSplitter::new(batch_size) - } else { - NoopBatchTransformer::new() //todo: figure this out lol - }, + batch_size: enforce_batch_size_in_joins.then(|| batch_size), }; let stream = async_try_stream(|mut emitter| async move { state.start_join_time(); - let result = state.join(&mut emitter).await; + let result = state.join(&context, &mut emitter).await; state.stop_join_time(); result }); @@ -565,7 +561,7 @@ fn stats_cartesian_product( } /// A stream that issues [RecordBatch]es as they arrive from the right of the join. -struct CrossJoinStream { +struct CrossJoinStream { /// Input schema schema: Arc, /// Future for data from left side @@ -580,8 +576,8 @@ struct CrossJoinStream { processed_right_data: Option, /// Left data (copy of the entire buffered left side) left_data: RecordBatch, - /// Batch transformer - batch_transformer: T, + /// Max batch size + batch_size: Option, } fn build_batch( @@ -612,8 +608,11 @@ fn build_batch( .map_err(Into::into) } -impl CrossJoinStream { - async fn join(&mut self, cx: &mut std::task::Context<'_>) -> Result<()> { +impl CrossJoinStream { + async fn join( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { if !self.collect_build_side(cx)? { return Ok(()); } @@ -623,13 +622,13 @@ impl CrossJoinStream { loop { return match self.state { CrossJoinStreamState::WaitBuildSide => { - handle_state!(ready!(self.collect_build_side(cx))) + handle_state!(ready!(self.collect_build_side())) } CrossJoinStreamState::FetchProbeBatch => { - handle_state!(ready!(self.fetch_probe_batch(cx))) + handle_state!(ready!(self.fetch_probe_batch())) } CrossJoinStreamState::BuildBatches(_) => { - let poll = handle_state!(self.build_batches()); + let poll = handle_state!(self.build_batches(emitter)); self.join_metrics.baseline.record_poll(poll) } }; @@ -639,7 +638,7 @@ impl CrossJoinStream { /// Collects build (left) side of the join into the state. In case of an empty build batch, /// the execution terminates. Otherwise, the state is updated to fetch probe (right) batch. /// Returns true if build side was loaded and non-empty - fn collect_build_side(&mut self, cx: &mut std::task::Context<'_>) -> Result { + fn collect_build_side(&mut self) -> Result { let build_timer = self.join_metrics.build_time.timer(); let left_data = match self.left_fut.get(cx) { //todo: figure out how to make this async maybe @@ -685,7 +684,8 @@ impl CrossJoinStream { /// If all the results are produced, the state is set to fetch new probe batch. async fn build_batches( &mut self, - ) -> Result>> { + emitter: &mut TryEmitter, + ) -> Result<()> { let right_batch = self .processed_right_data .ok_or(internal_err!("Expected RecordBatch for right side"))?; @@ -708,6 +708,7 @@ impl CrossJoinStream { self.left_index += 1; } + emitter.emit(batch); return Ok(StatefulStreamResult::Ready(Some(batch))); } } From 6ced5436288112024c82dd2ce77a0d8568595e04 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Tue, 11 Aug 2026 19:32:36 -0400 Subject: [PATCH 07/13] moved join function to use new async methods --- .../physical-plan/src/joins/cross_join.rs | 66 +++++++++---------- 1 file changed, 32 insertions(+), 34 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 30024f719852e..b2793ad4cebd0 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -18,6 +18,7 @@ //! Defines the cross join plan for loading the left side of the cross join //! and producing batches in parallel for the right partitions +use std::future::poll_fn; use std::{sync::Arc, task::Poll}; use super::utils::{ @@ -360,12 +361,13 @@ impl ExecutionPlan for CrossJoinExec { join_metrics, processed_right_data: None, left_data: RecordBatch::new_empty(self.left().schema()), - batch_size: enforce_batch_size_in_joins.then(|| batch_size), + batch_size: enforce_batch_size_in_joins.then_some(batch_size), + batch_transformer: NoopBatchTransformer::new(), //todo fix, }; let stream = async_try_stream(|mut emitter| async move { state.start_join_time(); - let result = state.join(&context, &mut emitter).await; + let result = state.join(&mut emitter).await; state.stop_join_time(); result }); @@ -561,7 +563,7 @@ fn stats_cartesian_product( } /// A stream that issues [RecordBatch]es as they arrive from the right of the join. -struct CrossJoinStream { +struct CrossJoinStream { /// Input schema schema: Arc, /// Future for data from left side @@ -576,8 +578,10 @@ struct CrossJoinStream { processed_right_data: Option, /// Left data (copy of the entire buffered left side) left_data: RecordBatch, - /// Max batch size + /// Max batch size (todo: maybe remove) batch_size: Option, + // todo: add description if we're keeping it + batch_transformer: T, } fn build_batch( @@ -608,40 +612,35 @@ fn build_batch( .map_err(Into::into) } -impl CrossJoinStream { +impl CrossJoinStream +where + T: BatchTransformer, +{ async fn join( &mut self, emitter: &mut TryEmitter, ) -> Result<()> { - if !self.collect_build_side(cx)? { + if !self.collect_build_side().await? { return Ok(()); } - while self.processed_right_data.is_none() {} // <- maybe add helper functions for when we have more data to determine the looping condition? - - loop { - return match self.state { - CrossJoinStreamState::WaitBuildSide => { - handle_state!(ready!(self.collect_build_side())) - } - CrossJoinStreamState::FetchProbeBatch => { - handle_state!(ready!(self.fetch_probe_batch())) - } - CrossJoinStreamState::BuildBatches(_) => { - let poll = handle_state!(self.build_batches(emitter)); - self.join_metrics.baseline.record_poll(poll) + while self.fetch_probe_batch().await? { + loop { + if !self.build_batches(emitter).await? { + break; } - }; + } } + + Ok(()) } /// Collects build (left) side of the join into the state. In case of an empty build batch, /// the execution terminates. Otherwise, the state is updated to fetch probe (right) batch. /// Returns true if build side was loaded and non-empty - fn collect_build_side(&mut self) -> Result { + async fn collect_build_side(&mut self) -> Result { let build_timer = self.join_metrics.build_time.timer(); - let left_data = match self.left_fut.get(cx) { - //todo: figure out how to make this async maybe + let left_data = match poll_fn(|cx| self.left_fut.get(cx)).await { Ok(left_data) => left_data, Err(e) => return Err(e), }; @@ -659,9 +658,7 @@ impl CrossJoinStream { /// Fetches the probe (right) batch, updates the metrics, and save the batch in the state. /// Then, the state is updated to build result batches. - async fn fetch_probe_batch( - &mut self, - ) -> Result>> { + async fn fetch_probe_batch(&mut self) -> Result { self.left_index = 0; let right_data = match self.right.next().await { Some(Ok(right_data)) => right_data, @@ -670,22 +667,25 @@ impl CrossJoinStream { // Release the right (probe) input pipeline's resources. let right_schema = self.right.schema(); self.right = Box::pin(EmptyRecordBatchStream::new(right_schema)); - return Ok(StatefulStreamResult::Ready(None)); + return Ok(false); } }; self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(right_data.num_rows()); self.processed_right_data = Some(right_data); - Ok(StatefulStreamResult::Continue) + Ok(true) } /// Joins the indexed row of left data with the current probe batch. /// If all the results are produced, the state is set to fetch new probe batch. + /// Err -> error lol + /// true -> emitted a batch, keep going + /// false -> fetch more from the probe async fn build_batches( &mut self, emitter: &mut TryEmitter, - ) -> Result<()> { + ) -> Result { let right_batch = self .processed_right_data .ok_or(internal_err!("Expected RecordBatch for right side"))?; @@ -708,14 +708,12 @@ impl CrossJoinStream { self.left_index += 1; } - emitter.emit(batch); - return Ok(StatefulStreamResult::Ready(Some(batch))); + emitter.emit(batch).await; + return Ok(true); } } - } else { - self.state = CrossJoinStreamState::FetchProbeBatch; } - Ok(StatefulStreamResult::Continue) + Ok(false) } } From 91c9c0f96796303b327552f03e712528c27d4407 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Tue, 11 Aug 2026 19:53:26 -0400 Subject: [PATCH 08/13] handling batch sizing --- .../physical-plan/src/joins/cross_join.rs | 59 ++++++------------- 1 file changed, 19 insertions(+), 40 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index b2793ad4cebd0..b5ee3e5e9d93d 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -362,7 +362,6 @@ impl ExecutionPlan for CrossJoinExec { processed_right_data: None, left_data: RecordBatch::new_empty(self.left().schema()), batch_size: enforce_batch_size_in_joins.then_some(batch_size), - batch_transformer: NoopBatchTransformer::new(), //todo fix, }; let stream = async_try_stream(|mut emitter| async move { @@ -563,25 +562,21 @@ fn stats_cartesian_product( } /// A stream that issues [RecordBatch]es as they arrive from the right of the join. -struct CrossJoinStream { +struct CrossJoinStream { /// Input schema schema: Arc, /// Future for data from left side left_fut: OnceFut, /// Right side stream right: SendableRecordBatchStream, - /// Current value on the left - left_index: usize, /// Join execution metrics join_metrics: BuildProbeJoinMetrics, /// Currently processed right side batch processed_right_data: Option, /// Left data (copy of the entire buffered left side) left_data: RecordBatch, - /// Max batch size (todo: maybe remove) + /// Max batch size batch_size: Option, - // todo: add description if we're keeping it - batch_transformer: T, } fn build_batch( @@ -612,10 +607,7 @@ fn build_batch( .map_err(Into::into) } -impl CrossJoinStream -where - T: BatchTransformer, -{ +impl CrossJoinStream { async fn join( &mut self, emitter: &mut TryEmitter, @@ -659,7 +651,6 @@ where /// Fetches the probe (right) batch, updates the metrics, and save the batch in the state. /// Then, the state is updated to build result batches. async fn fetch_probe_batch(&mut self) -> Result { - self.left_index = 0; let right_data = match self.right.next().await { Some(Ok(right_data)) => right_data, Some(Err(e)) => return Err(e), @@ -682,38 +673,26 @@ where /// Err -> error lol /// true -> emitted a batch, keep going /// false -> fetch more from the probe - async fn build_batches( + async fn process_right_batch( &mut self, + right_batch: &RecordBatch, emitter: &mut TryEmitter, - ) -> Result { - let right_batch = self - .processed_right_data - .ok_or(internal_err!("Expected RecordBatch for right side"))?; - if self.left_index < self.left_data.num_rows() { - match self.batch_transformer.next() { - None => { - let join_timer = self.join_metrics.join_time.timer(); - let result = build_batch( - self.left_index, - &right_batch, - &self.left_data, - &self.schema, - ); - join_timer.done(); - - self.batch_transformer.set_batch(result?); - } - Some((batch, last)) => { - if last { - self.left_index += 1; - } - - emitter.emit(batch).await; - return Ok(true); - } + ) -> Result<()> { + for left_index in 0..self.left_data.num_rows() { + let join_timer = self.join_metrics.join_time.timer(); + let result = + build_batch(left_index, right_batch, &self.left_data, &self.schema)?; + join_timer.done(); + + // TODO: understand the batching stuff - do we only want one batch, then fetch more from the right, then continue here? Or what + if let Some(batch_size) = self.batch_size { + let x = result.slice(offset, length); + } else { + emitter.emit(result).await; } } - Ok(false) + + Ok(()) } } From 98b99f4d5b63c76ebc19a14989e7a1aa0b36074b Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Wed, 12 Aug 2026 11:47:41 -0400 Subject: [PATCH 09/13] adding comments --- .../physical-plan/src/joins/cross_join.rs | 44 ++++++++----------- 1 file changed, 19 insertions(+), 25 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index b5ee3e5e9d93d..b1a3ca66e5da7 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -18,6 +18,7 @@ //! Defines the cross join plan for loading the left side of the cross join //! and producing batches in parallel for the right partitions +use std::cmp::min; use std::future::poll_fn; use std::{sync::Arc, task::Poll}; @@ -42,6 +43,7 @@ use crate::{ use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::compute::concat_batches; +use arrow::compute::kernels::length; use arrow::datatypes::{Fields, Schema, SchemaRef}; use datafusion_common::stats::Precision; use datafusion_common::{ @@ -54,6 +56,7 @@ use datafusion_physical_expr::equivalence::join_equivalence_properties; use async_trait::async_trait; use futures::{Stream, StreamExt, TryStreamExt, ready}; +use log::Record; /// Data of the left side that is buffered into memory #[derive(Debug)] @@ -357,9 +360,7 @@ impl ExecutionPlan for CrossJoinExec { schema: Arc::clone(&self.schema), left_fut, right: stream, - left_index: 0, join_metrics, - processed_right_data: None, left_data: RecordBatch::new_empty(self.left().schema()), batch_size: enforce_batch_size_in_joins.then_some(batch_size), }; @@ -571,8 +572,6 @@ struct CrossJoinStream { right: SendableRecordBatchStream, /// Join execution metrics join_metrics: BuildProbeJoinMetrics, - /// Currently processed right side batch - processed_right_data: Option, /// Left data (copy of the entire buffered left side) left_data: RecordBatch, /// Max batch size @@ -608,6 +607,7 @@ fn build_batch( } impl CrossJoinStream { + // Collect the left (build) side, then continue processing the right side against it until we have no more rows on the right async fn join( &mut self, emitter: &mut TryEmitter, @@ -616,19 +616,14 @@ impl CrossJoinStream { return Ok(()); } - while self.fetch_probe_batch().await? { - loop { - if !self.build_batches(emitter).await? { - break; - } - } + while let Some(right_batch) = self.fetch_probe_batch().await? { + self.process_right_batch(&right_batch, emitter).await? } Ok(()) } - /// Collects build (left) side of the join into the state. In case of an empty build batch, - /// the execution terminates. Otherwise, the state is updated to fetch probe (right) batch. + /// Collects build (left) side of the join into the state. In case of an empty build batch, the execution terminates. /// Returns true if build side was loaded and non-empty async fn collect_build_side(&mut self) -> Result { let build_timer = self.join_metrics.build_time.timer(); @@ -648,9 +643,8 @@ impl CrossJoinStream { Ok(result) } - /// Fetches the probe (right) batch, updates the metrics, and save the batch in the state. - /// Then, the state is updated to build result batches. - async fn fetch_probe_batch(&mut self) -> Result { + /// Fetches the probe (right) batch, updates the metrics, and returns the batch + async fn fetch_probe_batch(&mut self) -> Result> { let right_data = match self.right.next().await { Some(Ok(right_data)) => right_data, Some(Err(e)) => return Err(e), @@ -658,21 +652,16 @@ impl CrossJoinStream { // Release the right (probe) input pipeline's resources. let right_schema = self.right.schema(); self.right = Box::pin(EmptyRecordBatchStream::new(right_schema)); - return Ok(false); + return Ok(None); } }; self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(right_data.num_rows()); - self.processed_right_data = Some(right_data); - Ok(true) + Ok(Some(right_data)) } - /// Joins the indexed row of left data with the current probe batch. - /// If all the results are produced, the state is set to fetch new probe batch. - /// Err -> error lol - /// true -> emitted a batch, keep going - /// false -> fetch more from the probe + /// Joins the left data with the current probe batch, using the emitter to emit the resultant batches async fn process_right_batch( &mut self, right_batch: &RecordBatch, @@ -684,9 +673,14 @@ impl CrossJoinStream { build_batch(left_index, right_batch, &self.left_data, &self.schema)?; join_timer.done(); - // TODO: understand the batching stuff - do we only want one batch, then fetch more from the right, then continue here? Or what if let Some(batch_size) = self.batch_size { - let x = result.slice(offset, length); + let mut offset = 0; + while offset < result.num_rows() { + let length = min(result.num_rows() - offset, batch_size); + let sliced_result = result.slice(offset, length); + emitter.emit(sliced_result).await; + offset += length; + } } else { emitter.emit(result).await; } From b4ac3e0efbcaa8866365609dbca3f177c7f6aac3 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Wed, 12 Aug 2026 12:08:28 -0400 Subject: [PATCH 10/13] added tests --- .../physical-plan/src/joins/cross_join.rs | 258 +++++++++++++++--- 1 file changed, 221 insertions(+), 37 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index b1a3ca66e5da7..79e2ce1c33994 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -20,11 +20,10 @@ use std::cmp::min; use std::future::poll_fn; -use std::{sync::Arc, task::Poll}; +use std::sync::Arc; use super::utils::{ - BatchSplitter, BatchTransformer, BuildProbeJoinMetrics, NoopBatchTransformer, - OnceAsync, OnceFut, StatefulStreamResult, adjust_right_output_partitioning, + BuildProbeJoinMetrics, OnceAsync, OnceFut, adjust_right_output_partitioning, reorder_output_after_swap, }; use crate::execution_plan::{EmissionType, boundedness_from_children}; @@ -37,26 +36,23 @@ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{ ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, check_if_same_properties, handle_state, + ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, Statistics, + check_if_same_properties, }; use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::compute::concat_batches; -use arrow::compute::kernels::length; use arrow::datatypes::{Fields, Schema, SchemaRef}; use datafusion_common::stats::Precision; use datafusion_common::{ DataFusionError, JoinType, Result, ScalarValue, assert_eq_or_internal_err, - internal_err, }; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use datafusion_physical_expr::equivalence::join_equivalence_properties; -use async_trait::async_trait; -use futures::{Stream, StreamExt, TryStreamExt, ready}; -use log::Record; +use futures::{StreamExt, TryStreamExt}; +use num_traits::Zero; /// Data of the left side that is buffered into memory #[derive(Debug)] @@ -210,7 +206,7 @@ async fn load_left_input( let left_schema = stream.schema(); // Load all batches and count the rows - let (batches, _metrics, reservation) = stream + let (batches, metrics, reservation) = stream .try_fold( (Vec::new(), metrics, reservation), |(mut batches, metrics, reservation), batch| async { @@ -228,7 +224,9 @@ async fn load_left_input( ) .await?; + let build_timer = metrics.build_time.timer(); let merged_batch = concat_batches(&left_schema, &batches)?; + build_timer.done(); Ok(JoinLeftData { merged_batch, @@ -365,17 +363,14 @@ impl ExecutionPlan for CrossJoinExec { batch_size: enforce_batch_size_in_joins.then_some(batch_size), }; - let stream = async_try_stream(|mut emitter| async move { - state.start_join_time(); - let result = state.join(&mut emitter).await; - state.stop_join_time(); - result - }); - // ObservedStream records the baseline metrics (output rows/batches, - // end time) + let schema = Arc::clone(&self.schema); + let baseline_metrics = state.join_metrics.baseline.clone(); + let stream = + async_try_stream(|mut emitter| async move { state.join(&mut emitter).await }); + Ok(Box::pin(ObservedStream::new( - Box::pin(RecordBatchStreamAdapter::new(state.schema, stream)), - baseline_metrics, //todo: fix this too + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + baseline_metrics, None, ))) } @@ -626,21 +621,16 @@ impl CrossJoinStream { /// Collects build (left) side of the join into the state. In case of an empty build batch, the execution terminates. /// Returns true if build side was loaded and non-empty async fn collect_build_side(&mut self) -> Result { - let build_timer = self.join_metrics.build_time.timer(); - let left_data = match poll_fn(|cx| self.left_fut.get(cx)).await { - Ok(left_data) => left_data, - Err(e) => return Err(e), - }; - build_timer.done(); - - let left_data = left_data.merged_batch.clone(); - let result = if left_data.num_rows() == 0 { - false - } else { - self.left_data = left_data; - true - }; - Ok(result) + let left_data = poll_fn(|cx| { + self.left_fut + .get(cx) + .map(|res| res.map(|data| data.merged_batch.clone())) + }) + .await?; + + let is_empty = left_data.num_rows().is_zero(); + self.left_data = left_data; + Ok(!is_empty) } /// Fetches the probe (right) batch, updates the metrics, and returns the batch @@ -694,9 +684,13 @@ impl CrossJoinStream { mod tests { use super::*; use crate::common; - use crate::test::{assert_join_metrics, build_table_scan_i32}; + use crate::test::{assert_join_metrics, build_table_i32, build_table_scan_i32}; + use std::time::Duration; + + use datafusion_common::instant::Instant; use datafusion_common::{assert_contains, test_util::batches_to_sort_string}; + use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use insta::assert_snapshot; @@ -987,6 +981,196 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_join_enforce_batch_size_splits_output() -> Result<()> { + let mut config = SessionConfig::new().with_batch_size(2); + config.options_mut().execution.enforce_batch_size_in_joins = true; + let task_ctx = Arc::new(TaskContext::default().with_session_config(config)); + + let left = build_table_scan_i32( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 6]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table_scan_i32( + ("a2", &vec![10, 11, 12, 13, 14]), + ("b2", &vec![15, 16, 17, 18, 19]), + ("c2", &vec![20, 21, 22, 23, 24]), + ); + + let (_, batches, _) = join_collect(left, right, task_ctx).await?; + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 15); + assert!(batches.iter().all(|b| b.num_rows() <= 2)); + + Ok(()) + } + + fn delayed_stream( + batches: Vec, + delay: Duration, + ) -> SendableRecordBatchStream { + let schema = batches[0].schema(); + Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches.into_iter().map(Ok)).then( + move |item| async move { + tokio::time::sleep(delay).await; + item + }, + ), + )) + } + + fn probe_batches() -> Vec { + vec![ + build_table_i32( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ), + build_table_i32( + ("a2", &vec![20, 21]), + ("b2", &vec![22, 23]), + ("c2", &vec![24, 25]), + ), + build_table_i32( + ("a2", &vec![30, 31]), + ("b2", &vec![32, 33]), + ("c2", &vec![34, 35]), + ), + ] + } + + fn cross_join_stream( + left_batch: RecordBatch, + right: SendableRecordBatchStream, + batch_size: Option, + ) -> Result<(SendableRecordBatchStream, ExecutionPlanMetricsSet)> { + let metrics = ExecutionPlanMetricsSet::new(); + let join_metrics = BuildProbeJoinMetrics::new(0, &metrics); + + let runtime = RuntimeEnvBuilder::new().build_arc()?; + let reservation = MemoryConsumer::new("test").register(&runtime.memory_pool); + + let left_schema = left_batch.schema(); + let mut fields: Vec<_> = left_schema.fields().iter().cloned().collect(); + fields.extend(right.schema().fields().iter().cloned()); + let schema = Arc::new(Schema::new(fields)); + + let left_data = JoinLeftData { + merged_batch: left_batch, + _reservation: reservation, + }; + let left_fut = OnceFut::new(async move { Ok(left_data) }); + + let mut state = CrossJoinStream { + schema: Arc::clone(&schema), + left_fut, + right, + join_metrics, + left_data: RecordBatch::new_empty(left_schema), + batch_size, + }; + let baseline_metrics = state.join_metrics.baseline.clone(); + let stream = + async_try_stream( + move |mut emitter| async move { state.join(&mut emitter).await }, + ); + let observed = ObservedStream::new( + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + baseline_metrics, + None, + ); + Ok((Box::pin(observed), metrics)) + } + + fn elapsed_compute_of(metrics: &ExecutionPlanMetricsSet) -> Duration { + Duration::from_nanos(metrics.clone_inner().elapsed_compute().unwrap_or(0) as u64) + } + + async fn check_elapsed_compute_excluded(mut run: F) -> Result<()> + where + F: FnMut(Duration) -> Fut, + Fut: Future>, + { + let mut delay = Duration::from_millis(50); + for attempt in 0..3 { + let (elapsed_compute, wall) = run(delay).await?; + if elapsed_compute < delay { + return Ok(()); + } + assert!( + attempt < 2, + "elapsed_compute ({elapsed_compute:?}) should be well below the \ + injected delay ({delay:?}); wall {wall:?}" + ); + delay *= 4; + } + unreachable!() + } + + #[tokio::test] + async fn elapsed_compute_excludes_probe_input_wait() -> Result<()> { + check_elapsed_compute_excluded(|delay| async move { + let left = build_table_i32( + ("a1", &vec![1, 2]), + ("b1", &vec![3, 4]), + ("c1", &vec![5, 6]), + ); + let right = delayed_stream(probe_batches(), delay); + let (stream, metrics) = cross_join_stream(left, right, None)?; + + let start = Instant::now(); + let batches = common::collect(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 12); + assert!( + wall >= delay * 3, + "probe delays should dominate wall time, got {wall:?}" + ); + Ok((elapsed_compute_of(&metrics), wall)) + }) + .await + } + + #[tokio::test] + async fn elapsed_compute_excludes_consumer_wait() -> Result<()> { + check_elapsed_compute_excluded(|delay| async move { + let left = build_table_i32( + ("a1", &vec![1, 2]), + ("b1", &vec![3, 4]), + ("c1", &vec![5, 6]), + ); + let right = delayed_stream(probe_batches(), Duration::ZERO); + let (mut stream, metrics) = cross_join_stream(left, right, Some(1))?; + + let start = Instant::now(); + let mut output_batches = 0u32; + while let Some(batch) = stream.next().await { + batch?; + output_batches += 1; + tokio::time::sleep(delay).await; + } + drop(stream); + let wall = start.elapsed(); + + assert!( + output_batches >= 3, + "expected multiple emitted batches, got {output_batches}" + ); + assert!( + wall >= delay * output_batches, + "consumer delays should dominate wall time, got {wall:?}" + ); + Ok((elapsed_compute_of(&metrics), wall)) + }) + .await + } + /// Returns the column names on the schema fn columns(schema: &Schema) -> Vec { schema.fields().iter().map(|f| f.name().clone()).collect() From 59c4df0d618d56fa356828f86f031013aee16f1a Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Thu, 13 Aug 2026 12:06:21 -0400 Subject: [PATCH 11/13] removing batch sizing --- .../physical-plan/src/joins/cross_join.rs | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 9d60c88572a11..445e6233b15de 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -350,10 +350,6 @@ impl ExecutionPlan for CrossJoinExec { let reservation = MemoryConsumer::new("CrossJoinExec").register(context.memory_pool()); - let batch_size = context.session_config().batch_size(); - let enforce_batch_size_in_joins = - context.session_config().enforce_batch_size_in_joins(); - let left_fut = self.left_fut.try_once(|| { let left_stream = self.left.execute(0, context)?; @@ -370,7 +366,6 @@ impl ExecutionPlan for CrossJoinExec { right: stream, join_metrics, left_data: RecordBatch::new_empty(self.left().schema()), - batch_size: enforce_batch_size_in_joins.then_some(batch_size), }; let schema = Arc::clone(&self.schema); @@ -579,8 +574,6 @@ struct CrossJoinStream { join_metrics: BuildProbeJoinMetrics, /// Left data (copy of the entire buffered left side) left_data: RecordBatch, - /// Max batch size - batch_size: Option, } fn build_batch( @@ -673,17 +666,7 @@ impl CrossJoinStream { build_batch(left_index, right_batch, &self.left_data, &self.schema)?; join_timer.done(); - if let Some(batch_size) = self.batch_size { - let mut offset = 0; - while offset < result.num_rows() { - let length = min(result.num_rows() - offset, batch_size); - let sliced_result = result.slice(offset, length); - emitter.emit(sliced_result).await; - offset += length; - } - } else { - emitter.emit(result).await; - } + emitter.emit(result).await; } Ok(()) From 9a7ad53a77ac8831423886d2ed11aa074a866f1f Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Thu, 13 Aug 2026 12:26:10 -0400 Subject: [PATCH 12/13] addressing PR comments --- .../physical-plan/src/joins/cross_join.rs | 216 +----------------- 1 file changed, 10 insertions(+), 206 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 445e6233b15de..4d09075eaccae 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -18,7 +18,6 @@ //! Defines the cross join plan for loading the left side of the cross join //! and producing batches in parallel for the right partitions -use std::cmp::min; use std::future::poll_fn; use std::sync::Arc; @@ -614,9 +613,7 @@ impl CrossJoinStream { return Ok(()); } - while let Some(right_batch) = self.fetch_probe_batch().await? { - self.process_right_batch(&right_batch, emitter).await? - } + self.process_right_batch(emitter).await?; Ok(()) } @@ -657,16 +654,17 @@ impl CrossJoinStream { /// Joins the left data with the current probe batch, using the emitter to emit the resultant batches async fn process_right_batch( &mut self, - right_batch: &RecordBatch, emitter: &mut TryEmitter, ) -> Result<()> { - for left_index in 0..self.left_data.num_rows() { - let join_timer = self.join_metrics.join_time.timer(); - let result = - build_batch(left_index, right_batch, &self.left_data, &self.schema)?; - join_timer.done(); + while let Some(right_batch) = self.fetch_probe_batch().await? { + for left_index in 0..self.left_data.num_rows() { + let join_timer = self.join_metrics.join_time.timer(); + let result = + build_batch(left_index, &right_batch, &self.left_data, &self.schema)?; + join_timer.done(); - emitter.emit(result).await; + emitter.emit(result).await; + } } Ok(()) @@ -677,13 +675,9 @@ impl CrossJoinStream { mod tests { use super::*; use crate::common; - use crate::test::{assert_join_metrics, build_table_i32, build_table_scan_i32}; + use crate::test::{assert_join_metrics, build_table_scan_i32}; - use std::time::Duration; - - use datafusion_common::instant::Instant; use datafusion_common::{assert_contains, test_util::batches_to_sort_string}; - use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use insta::assert_snapshot; @@ -974,196 +968,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_join_enforce_batch_size_splits_output() -> Result<()> { - let mut config = SessionConfig::new().with_batch_size(2); - config.options_mut().execution.enforce_batch_size_in_joins = true; - let task_ctx = Arc::new(TaskContext::default().with_session_config(config)); - - let left = build_table_scan_i32( - ("a1", &vec![1, 2, 3]), - ("b1", &vec![4, 5, 6]), - ("c1", &vec![7, 8, 9]), - ); - let right = build_table_scan_i32( - ("a2", &vec![10, 11, 12, 13, 14]), - ("b2", &vec![15, 16, 17, 18, 19]), - ("c2", &vec![20, 21, 22, 23, 24]), - ); - - let (_, batches, _) = join_collect(left, right, task_ctx).await?; - - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, 15); - assert!(batches.iter().all(|b| b.num_rows() <= 2)); - - Ok(()) - } - - fn delayed_stream( - batches: Vec, - delay: Duration, - ) -> SendableRecordBatchStream { - let schema = batches[0].schema(); - Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::iter(batches.into_iter().map(Ok)).then( - move |item| async move { - tokio::time::sleep(delay).await; - item - }, - ), - )) - } - - fn probe_batches() -> Vec { - vec![ - build_table_i32( - ("a2", &vec![10, 11]), - ("b2", &vec![12, 13]), - ("c2", &vec![14, 15]), - ), - build_table_i32( - ("a2", &vec![20, 21]), - ("b2", &vec![22, 23]), - ("c2", &vec![24, 25]), - ), - build_table_i32( - ("a2", &vec![30, 31]), - ("b2", &vec![32, 33]), - ("c2", &vec![34, 35]), - ), - ] - } - - fn cross_join_stream( - left_batch: RecordBatch, - right: SendableRecordBatchStream, - batch_size: Option, - ) -> Result<(SendableRecordBatchStream, ExecutionPlanMetricsSet)> { - let metrics = ExecutionPlanMetricsSet::new(); - let join_metrics = BuildProbeJoinMetrics::new(0, &metrics); - - let runtime = RuntimeEnvBuilder::new().build_arc()?; - let reservation = MemoryConsumer::new("test").register(&runtime.memory_pool); - - let left_schema = left_batch.schema(); - let mut fields: Vec<_> = left_schema.fields().iter().cloned().collect(); - fields.extend(right.schema().fields().iter().cloned()); - let schema = Arc::new(Schema::new(fields)); - - let left_data = JoinLeftData { - merged_batch: left_batch, - _reservation: reservation, - }; - let left_fut = OnceFut::new(async move { Ok(left_data) }); - - let mut state = CrossJoinStream { - schema: Arc::clone(&schema), - left_fut, - right, - join_metrics, - left_data: RecordBatch::new_empty(left_schema), - batch_size, - }; - let baseline_metrics = state.join_metrics.baseline.clone(); - let stream = - async_try_stream( - move |mut emitter| async move { state.join(&mut emitter).await }, - ); - let observed = ObservedStream::new( - Box::pin(RecordBatchStreamAdapter::new(schema, stream)), - baseline_metrics, - None, - ); - Ok((Box::pin(observed), metrics)) - } - - fn elapsed_compute_of(metrics: &ExecutionPlanMetricsSet) -> Duration { - Duration::from_nanos(metrics.clone_inner().elapsed_compute().unwrap_or(0) as u64) - } - - async fn check_elapsed_compute_excluded(mut run: F) -> Result<()> - where - F: FnMut(Duration) -> Fut, - Fut: Future>, - { - let mut delay = Duration::from_millis(50); - for attempt in 0..3 { - let (elapsed_compute, wall) = run(delay).await?; - if elapsed_compute < delay { - return Ok(()); - } - assert!( - attempt < 2, - "elapsed_compute ({elapsed_compute:?}) should be well below the \ - injected delay ({delay:?}); wall {wall:?}" - ); - delay *= 4; - } - unreachable!() - } - - #[tokio::test] - async fn elapsed_compute_excludes_probe_input_wait() -> Result<()> { - check_elapsed_compute_excluded(|delay| async move { - let left = build_table_i32( - ("a1", &vec![1, 2]), - ("b1", &vec![3, 4]), - ("c1", &vec![5, 6]), - ); - let right = delayed_stream(probe_batches(), delay); - let (stream, metrics) = cross_join_stream(left, right, None)?; - - let start = Instant::now(); - let batches = common::collect(stream).await?; - let wall = start.elapsed(); - - let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(rows, 12); - assert!( - wall >= delay * 3, - "probe delays should dominate wall time, got {wall:?}" - ); - Ok((elapsed_compute_of(&metrics), wall)) - }) - .await - } - - #[tokio::test] - async fn elapsed_compute_excludes_consumer_wait() -> Result<()> { - check_elapsed_compute_excluded(|delay| async move { - let left = build_table_i32( - ("a1", &vec![1, 2]), - ("b1", &vec![3, 4]), - ("c1", &vec![5, 6]), - ); - let right = delayed_stream(probe_batches(), Duration::ZERO); - let (mut stream, metrics) = cross_join_stream(left, right, Some(1))?; - - let start = Instant::now(); - let mut output_batches = 0u32; - while let Some(batch) = stream.next().await { - batch?; - output_batches += 1; - tokio::time::sleep(delay).await; - } - drop(stream); - let wall = start.elapsed(); - - assert!( - output_batches >= 3, - "expected multiple emitted batches, got {output_batches}" - ); - assert!( - wall >= delay * output_batches, - "consumer delays should dominate wall time, got {wall:?}" - ); - Ok((elapsed_compute_of(&metrics), wall)) - }) - .await - } - /// Returns the column names on the schema fn columns(schema: &Schema) -> Vec { schema.fields().iter().map(|f| f.name().clone()).collect() From 0840347006358a17ba811bd0400ee0e248a94970 Mon Sep 17 00:00:00 2001 From: saadtajwar Date: Thu, 13 Aug 2026 12:43:32 -0400 Subject: [PATCH 13/13] lint --- datafusion/physical-plan/src/joins/cross_join.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index efde34ee70b3f..3e798573f5dec 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -34,10 +34,9 @@ use crate::projection::{ use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; use crate::{ - ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, - Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, - validate_child_count, + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, }; use arrow::array::{RecordBatch, RecordBatchOptions};