refactor: CrossJoinStream (simplifying, less state, async generator pattern) - #24291
refactor: CrossJoinStream (simplifying, less state, async generator pattern)#24291saadtajwar wants to merge 15 commits into
Conversation
saadtajwar
changed the title
…-crossjoinstream # Conflicts: # datafusion/physical-plan/src/joins/cross_join.rs
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24291 +/- ##
==========================================
- Coverage 81.14% 81.14% -0.01%
==========================================
Files 1112 1112
Lines 386933 386893 -40
Branches 386933 386893 -40
==========================================
- Hits 313978 313935 -43
- Misses 54462 54465 +3
Partials 18493 18493 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@rluvaton & @2010YOUY01 & @buraksenn - ready for review! Thanks in advance for the feedback - this was super fun to work on 😁 |
2010YOUY01
left a comment
2010YOUY01
left a comment
There was a problem hiding this comment.
Thank you! I think for the cross join's complexity, the explicit state management is not necessary, and the generator pattern make it look better.
I have left some suggestions for you to consider.
| right: stream, | ||
| join_metrics, | ||
| left_data: RecordBatch::new_empty(self.left().schema()), | ||
| batch_size: enforce_batch_size_in_joins.then_some(batch_size), |
There was a problem hiding this 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
There was a problem hiding this comment.
Ooh I see, yes that makes sense to me!
| } | ||
|
|
||
| while let Some(right_batch) = self.fetch_probe_batch().await? { | ||
| self.process_right_batch(&right_batch, emitter).await? |
There was a problem hiding this comment.
It seems we can put fetch_probe_batch inside process_right_batch to make it even simpler
There was a problem hiding this comment.
Ooh yeah - done!
| build_batch(left_index, right_batch, &self.left_data, &self.schema)?; | ||
| join_timer.done(); | ||
|
|
||
| if let Some(batch_size) = self.batch_size { |
There was a problem hiding this 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)
There was a problem hiding this comment.
Removed the batch sizing from the first comment :)
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_join_enforce_batch_size_splits_output() -> Result<()> { |
There was a problem hiding this comment.
| async fn test_join_enforce_batch_size_splits_output() -> Result<()> { | |
| // 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<()> { |
There was a problem hiding this comment.
Removed this test as we removed the batch sizing!
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn elapsed_compute_excludes_probe_input_wait() -> Result<()> { |
There was a problem hiding this 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.
There was a problem hiding this comment.
Removed these :)
|
@2010YOUY01 - thank you for the PR review! Pushed changes to address your comments, please let me know your thoughts! |
…-crossjoinstream # Conflicts: # datafusion/physical-plan/src/joins/cross_join.rs
3 participants
Which issue does this PR close?
Rationale for this change
The code for
CrossJoinStreamhad an opportunity to be simplified by removing much of the state related to the polling implementationWhat changes are included in this PR?
Removed the polling mechanisms in
CrossJoinStreamand the state associated with it, and instead simplified to a simple async-generator pattern of fetching the build side -> while there are batches available on the right, fetch them and perform the joinAre these changes tested?
Yes
Are there any user-facing changes?
No, these are internal to the physical execution of the join