You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "saadtajwar (via GitHub)" <gi...@apache.org> on 2026/08/12 16:12:36 UTC
[PR] refactor: CrossJoinStream (simplifying, less state, async) [datafusion]
saadtajwar opened a new pull request, #24291:
URL: https://github.com/apache/datafusion/pull/24291
## Which issue does this PR close?
- Part of #23974
## Rationale for this change
The code for `CrossJoinStream` had an opportunity to be simplified by removing much of the state related to the polling implementation
## What changes are included in this PR?
Removed the polling mechanisms in `CrossJoinStream` and the state associated with it, and instead simplified to a simple pattern of fetching the build side -> while there are batches available on the right, fetch them and perform the join
## Are these changes tested?
Yes
## Are there any user-facing changes?
No, these are internal to the physical execution of the join
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: CrossJoinStream (simplifying, less state, async generator pattern) [datafusion]
Posted by "saadtajwar (via GitHub)" <gi...@apache.org>.
saadtajwar commented on PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#issuecomment-5270173796
@rluvaton & @2010YOUY01 & @buraksenn - ready for review! Thanks in advance for the feedback - this was super fun to work on 😁
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: CrossJoinStream (simplifying, less state, async generator pattern) [datafusion]
Posted by "saadtajwar (via GitHub)" <gi...@apache.org>.
saadtajwar commented on code in PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#discussion_r3777175791
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -1054,6 +991,196 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_join_enforce_batch_size_splits_output() -> Result<()> {
Review Comment:
Removed this test as we removed the batch sizing!
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -361,29 +364,25 @@ 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,
- state: CrossJoinStreamState::WaitBuildSide,
- 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,
- state: CrossJoinStreamState::WaitBuildSide,
- 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,
+ join_metrics,
+ left_data: RecordBatch::new_empty(self.left().schema()),
+ batch_size: enforce_batch_size_in_joins.then_some(batch_size),
Review Comment:
Ooh I see, yes that makes sense to me!
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -641,129 +611,96 @@ fn build_batch(
.map_err(Into::into)
}
-#[async_trait]
-impl<T: BatchTransformer + Unpin + Send> Stream for CrossJoinStream<T> {
- type Item = Result<RecordBatch>;
-
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
- }
-}
-
-impl<T: BatchTransformer> CrossJoinStream<T> {
- /// Separate implementation function that unpins the [`CrossJoinStream`] so
- /// that partial borrows work correctly
- fn poll_next_impl(
+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,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- CrossJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- CrossJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- CrossJoinStreamState::BuildBatches(_) => {
- let poll = handle_state!(self.build_batches());
- self.join_metrics.baseline.record_poll(poll)
- }
- };
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ if !self.collect_build_side().await? {
+ return Ok(());
+ }
+
+ while let Some(right_batch) = self.fetch_probe_batch().await? {
+ self.process_right_batch(&right_batch, emitter).await?
Review Comment:
Ooh yeah - done!
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -1054,6 +991,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<RecordBatch>,
+ 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<RecordBatch> {
+ 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<usize>,
+ ) -> 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<F, Fut>(mut run: F) -> Result<()>
+ where
+ F: FnMut(Duration) -> Fut,
+ Fut: Future<Output = Result<(Duration, Duration)>>,
+ {
+ 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<()> {
Review Comment:
Removed these :)
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -641,129 +611,96 @@ fn build_batch(
.map_err(Into::into)
}
-#[async_trait]
-impl<T: BatchTransformer + Unpin + Send> Stream for CrossJoinStream<T> {
- type Item = Result<RecordBatch>;
-
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
- }
-}
-
-impl<T: BatchTransformer> CrossJoinStream<T> {
- /// Separate implementation function that unpins the [`CrossJoinStream`] so
- /// that partial borrows work correctly
- fn poll_next_impl(
+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,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- CrossJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- CrossJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- CrossJoinStreamState::BuildBatches(_) => {
- let poll = handle_state!(self.build_batches());
- self.join_metrics.baseline.record_poll(poll)
- }
- };
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ if !self.collect_build_side().await? {
+ return Ok(());
+ }
+
+ 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.
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- 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)),
- };
- build_timer.done();
-
- let left_data = left_data.merged_batch.clone();
- let result = if left_data.num_rows() == 0 {
- StatefulStreamResult::Ready(None)
- } else {
- self.left_data = left_data;
- self.state = CrossJoinStreamState::FetchProbeBatch;
- StatefulStreamResult::Continue
- };
- Poll::Ready(Ok(result))
+ /// 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<bool> {
+ 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 save the batch in the state.
- /// Then, the state is updated to build result batches.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- self.left_index = 0;
- let right_data = match ready!(self.right.poll_next_unpin(cx)) {
+ /// Fetches the probe (right) batch, updates the metrics, and returns the batch
+ async fn fetch_probe_batch(&mut self) -> Result<Option<RecordBatch>> {
+ 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(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))
+ 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.
- fn build_batches(&mut self) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let right_batch = self.state.try_as_record_batch()?;
- 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;
- }
-
- return Ok(StatefulStreamResult::Ready(Some(batch)));
+ /// 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<RecordBatch, DataFusionError>,
+ ) -> 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();
+
+ if let Some(batch_size) = self.batch_size {
Review Comment:
Removed the batch sizing from the first comment :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: CrossJoinStream (simplifying, less state, async generator pattern) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 merged PR #24291:
URL: https://github.com/apache/datafusion/pull/24291
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: CrossJoinStream (simplifying, less state, async generator pattern) [datafusion]
Posted by "saadtajwar (via GitHub)" <gi...@apache.org>.
saadtajwar commented on PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#issuecomment-5316673700
> Thanks again!
Thanks! I'll look to see if there are more streams to continue this work on, this was fun to work on!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: CrossJoinStream (simplifying, less state, async generator pattern) [datafusion]
Posted by "saadtajwar (via GitHub)" <gi...@apache.org>.
saadtajwar commented on PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#issuecomment-5283431444
@2010YOUY01 - thank you for the PR review! Pushed changes to address your comments, please let me know your thoughts!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: CrossJoinStream (simplifying, less state, async generator pattern) [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#issuecomment-5269752022
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/24291?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `92.72727%` with `16 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 81.30%. Comparing base ([`9b3b518`](https://app.codecov.io/gh/apache/datafusion/commit/9b3b518950c0c6a564b5525d6b274000f27dccd9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`5192985`](https://app.codecov.io/gh/apache/datafusion/commit/5192985015e3b9ade7fe6efadebefd24734fd01c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 1 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/24291?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [datafusion/physical-plan/src/joins/cross\_join.rs](https://app.codecov.io/gh/apache/datafusion/pull/24291?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Fjoins%2Fcross_join.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9qb2lucy9jcm9zc19qb2luLnJz) | 92.72% | [6 Missing and 10 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/24291?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #24291 +/- ##
========================================
Coverage 81.29% 81.30%
========================================
Files 1110 1110
Lines 385336 385440 +104
Branches 385336 385440 +104
========================================
+ Hits 313261 313370 +109
+ Misses 53594 53578 -16
- Partials 18481 18492 +11
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/24291?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] refactor: CrossJoinStream (simplifying, less state, async generator pattern) [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #24291:
URL: https://github.com/apache/datafusion/pull/24291#discussion_r3772272486
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -361,29 +364,25 @@ 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,
- state: CrossJoinStreamState::WaitBuildSide,
- 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,
- state: CrossJoinStreamState::WaitBuildSide,
- 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,
+ join_metrics,
+ left_data: RecordBatch::new_empty(self.left().schema()),
+ batch_size: enforce_batch_size_in_joins.then_some(batch_size),
Review Comment:
should we remove this `enforce_batch_size_in_joins`?
This should be introduced due to some legacy implementation tries to materialize `total_left_buffer_len * right_batch_len` at once, now it only materialize `one_left_row * right_batch_len` in each step, so the underlying restriction is already followed
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -641,129 +611,96 @@ fn build_batch(
.map_err(Into::into)
}
-#[async_trait]
-impl<T: BatchTransformer + Unpin + Send> Stream for CrossJoinStream<T> {
- type Item = Result<RecordBatch>;
-
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
- }
-}
-
-impl<T: BatchTransformer> CrossJoinStream<T> {
- /// Separate implementation function that unpins the [`CrossJoinStream`] so
- /// that partial borrows work correctly
- fn poll_next_impl(
+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,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- CrossJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- CrossJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- CrossJoinStreamState::BuildBatches(_) => {
- let poll = handle_state!(self.build_batches());
- self.join_metrics.baseline.record_poll(poll)
- }
- };
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ if !self.collect_build_side().await? {
+ return Ok(());
+ }
+
+ while let Some(right_batch) = self.fetch_probe_batch().await? {
+ self.process_right_batch(&right_batch, emitter).await?
Review Comment:
It seems we can put `fetch_probe_batch` inside `process_right_batch` to make it even simpler
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -641,129 +611,96 @@ fn build_batch(
.map_err(Into::into)
}
-#[async_trait]
-impl<T: BatchTransformer + Unpin + Send> Stream for CrossJoinStream<T> {
- type Item = Result<RecordBatch>;
-
- fn poll_next(
- mut self: std::pin::Pin<&mut Self>,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Self::Item>> {
- self.poll_next_impl(cx)
- }
-}
-
-impl<T: BatchTransformer> CrossJoinStream<T> {
- /// Separate implementation function that unpins the [`CrossJoinStream`] so
- /// that partial borrows work correctly
- fn poll_next_impl(
+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,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Option<Result<RecordBatch>>> {
- loop {
- return match self.state {
- CrossJoinStreamState::WaitBuildSide => {
- handle_state!(ready!(self.collect_build_side(cx)))
- }
- CrossJoinStreamState::FetchProbeBatch => {
- handle_state!(ready!(self.fetch_probe_batch(cx)))
- }
- CrossJoinStreamState::BuildBatches(_) => {
- let poll = handle_state!(self.build_batches());
- self.join_metrics.baseline.record_poll(poll)
- }
- };
+ emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
+ ) -> Result<()> {
+ if !self.collect_build_side().await? {
+ return Ok(());
+ }
+
+ 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.
- fn collect_build_side(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- 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)),
- };
- build_timer.done();
-
- let left_data = left_data.merged_batch.clone();
- let result = if left_data.num_rows() == 0 {
- StatefulStreamResult::Ready(None)
- } else {
- self.left_data = left_data;
- self.state = CrossJoinStreamState::FetchProbeBatch;
- StatefulStreamResult::Continue
- };
- Poll::Ready(Ok(result))
+ /// 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<bool> {
+ 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 save the batch in the state.
- /// Then, the state is updated to build result batches.
- fn fetch_probe_batch(
- &mut self,
- cx: &mut std::task::Context<'_>,
- ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> {
- self.left_index = 0;
- let right_data = match ready!(self.right.poll_next_unpin(cx)) {
+ /// Fetches the probe (right) batch, updates the metrics, and returns the batch
+ async fn fetch_probe_batch(&mut self) -> Result<Option<RecordBatch>> {
+ 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(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))
+ 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.
- fn build_batches(&mut self) -> Result<StatefulStreamResult<Option<RecordBatch>>> {
- let right_batch = self.state.try_as_record_batch()?;
- 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;
- }
-
- return Ok(StatefulStreamResult::Ready(Some(batch)));
+ /// 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<RecordBatch, DataFusionError>,
+ ) -> 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();
+
+ if let Some(batch_size) = self.batch_size {
Review Comment:
This slicing step seems not necessary. The `result` has length the same as `right_batch`, and the upstream operator should already ensure it's less than `batch_size` config (this is a global convention for operators)
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -1054,6 +991,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<RecordBatch>,
+ 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<RecordBatch> {
+ 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<usize>,
+ ) -> 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<F, Fut>(mut run: F) -> Result<()>
+ where
+ F: FnMut(Duration) -> Fut,
+ Fut: Future<Output = Result<(Duration, Duration)>>,
+ {
+ 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<()> {
Review Comment:
I found this test and the next one is quite hard to figure out their test goals, are they necessary? If they're auto generated we could remove them, else it would be great to better explain their test goals.
##########
datafusion/physical-plan/src/joins/cross_join.rs:
##########
@@ -1054,6 +991,196 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_join_enforce_batch_size_splits_output() -> Result<()> {
Review Comment:
```suggestion
// Cross join should follow the global output-size convention: all output batches
// should have fewer than `batch_size` rows, as specified in the config.
async fn test_cross_join_output_length() -> Result<()> {
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org