refactor: CrossJoinStream (simplifying, less state, async generator pattern) - #24291
refactor: CrossJoinStream (simplifying, less state, async generator pattern)#24291saadtajwar wants to merge 11 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.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 ☔ 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
| } | ||
|
|
||
| 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
| 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)
| } | ||
|
|
||
| #[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<()> { |
| } | ||
|
|
||
| #[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.
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