feat: add ASOF join physical operator - #23828
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #23828 +/- ##
==========================================
+ Coverage 81.19% 81.22% +0.03%
==========================================
Files 1110 1112 +2
Lines 388618 391419 +2801
Branches 388618 391419 +2801
==========================================
+ Hits 315531 317937 +2406
- Misses 54507 54794 +287
- Partials 18580 18688 +108 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2010YOUY01
left a comment
2010YOUY01
left a comment
There was a problem hiding this comment.
Thank you! This is a really good start. I have done a quick first pass and left some suggestions.
| vec![ChildStats::At(partition), ChildStats::Skip] | ||
| } | ||
|
|
||
| fn statistics_from_inputs( |
There was a problem hiding this comment.
Just an idea to make this PR smaller, could we use the default implementation here? We could implement it later in a follow-up PR.
There was a problem hiding this comment.
I kept this small override because ASOF has two exact facts the default would discard: the output row count equals the left row count, and unmodified left columns retain their statistics. Right-side column statistics remain unknown. I added a comment to make that scope explicit.
There was a problem hiding this comment.
I see, this makes sense.
There was a problem hiding this comment.
Thanks! I also updated it to preserve the same contract through the general projection.
|
There is no new commits after the previous review, you may have forgotten to push the local changes 🤔 @Xuanwo |
2010YOUY01
left a comment
2010YOUY01
left a comment
There was a problem hiding this comment.
I went over the implementation in detail, and I think it's well organized.
Need to do before merging
Before merging, I think we could add a few basic tests that run the executor and assert the results, just to cover some different cases:
- No equality condition in
on. - Different comparison operators in the match condition, such as
<and>=. - Complex expressions in the
onor match conditions, such asMATCH_CONDITION (l_c1 + l_c2) > r_c1orON l_c1 = (r_c1 + r_c2). I think these should be supported.
Optional suggestions
The main implementation/design questions I have are:
- We need to buffer
batch_sizeoutput rows before emitting them, now it's implemented through thePendingRowsstruct, so we need to buffer all source batches to produce the final output. I think this uses more memory and makes the implementation more complex. An alternative is a) keep ain_progress_batch, and emit it after it reaches threshould b) only buffer one (left_batch, right_batch), and materialize its valid indices into thein_progress_batchwhen the cursor moves across it. This might be fast enough and simpler. - Each loop iteration only advances the right index by one. We can probably explore some fast-forwarding optimization here.
But I suggest we first implement this end to end, including planning, SQL support, more tests, and benchmarks, before exploring these optimizations further. It should be easier to validate the ideas afterward.
For now, I only suggest trying to simplify the existing implementation or adding more documentation to make future iterations easier. I left a few suggestions in the comments.
|
|
||
| fn input_distribution_requirements(&self) -> InputDistributionRequirements { | ||
| InputDistributionRequirements::new(vec![ | ||
| Distribution::UnspecifiedDistribution, |
There was a problem hiding this comment.
Is it the case the optimizer will insert round-robin repartition, if we declare this UnspecifiedDistribution? It should be fine if it's doing so, I was a little bit confused by this name 🤔
| vec![ChildStats::At(partition), ChildStats::Skip] | ||
| } | ||
|
|
||
| fn statistics_from_inputs( |
There was a problem hiding this comment.
I see, this makes sense.
Co-authored-by: Yongting You <2010youy01@gmail.com>
|
Thanks for the detailed review — this was very helpful. I added result-based coverage for no equality keys, both |
|
Great! I got one tricky bug found by AI, otherwise it should be good to go. Found one blocking correctness issue.
I'm not so sure if there is an easy fix, it it's hard we could first reject floats in the planning, and fix that in a separate PR. |
Checked a bit, seems to be a mid level fix. Let's fix that in a seperate PR #24375. We can get this in first! 💌 |
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
I hope we could have ASOFJoinStream that handles the AsOfJoinStreamState and performs the actual join. We could take inspiration from the existing join implementations, such as HashJoinStream.
| #[derive(Default)] | ||
| struct PendingRows { | ||
| /// Distinct source batches referenced by `indices`. | ||
| sources: Vec<Arc<RecordBatch>>, |
There was a problem hiding this comment.
Do we need Vec<Arc<RecordBatch>> and not just Vec<RecordBatch>
There was a problem hiding this comment.
Yep, Arc is intentional here: it keeps per-row clones O(1) and gives PendingRows a stable identity for deduplication. Added a short comment.
| /// flush pending rows without resetting either cursor or the candidate | ||
| /// ``` | ||
| async fn next_batch(&mut self) -> Result<Option<RecordBatch>> { | ||
| fn poll_next_impl( |
There was a problem hiding this comment.
I haven't fully figure out what's the best practice to implement the state machines to make them easy to read and extend, probably it is:
- generator pattern [Status Update] Simplifing Streams to be more textbook-like and have less state while keeping the same perf #23974
- If we need to draw a state machine to understand its mechanism, then use explicit state management like
Not an issue for now, we could experiment in the future.
|
cc @jayzhan211, @xudong963, and @alamb, would you like to take another look and move forward? |
datafusion project typically wait for 24 hrs after approval to merge, so others can have a chance to look. This one is a larger PR, probably we could wait for 2 days, unless anyone need more time to review. https://datafusion.apache.org/contributor-guide/pr_review.html#pr-review-mechanics |
|
Got it, thank you @2010YOUY01 for the explanation 🙌 |
4 participants
Which issue does this PR close?
Rationale for this change
This is the first layer of the ASOF JOIN stack. It establishes a
broadcast-based physical execution contract independently so later
floating-point equality, logical-plan, SQL, DataFrame, and serialization changes
can be reviewed as smaller follow-up PRs.
The initial implementation deliberately favors the simpler broadcast design:
the right input must fit in memory and each left partition scans the shared
right-side batches. A repartitioned implementation can be evaluated separately
without changing the ASOF semantics introduced here.
Floating-point equality keys are rejected in this base layer because Arrow's
required sort order distinguishes
-0.0from+0.0while join equality doesnot. #24375 adds the required ordering normalization as an independently
reviewable layer.
What changes are included in this PR?
AsOfJoinExecfor left-preserving, Snowflake-style ASOF semantics.left partitions.
preserve the left-side output partitioning.
batches are zero-copy slices, and expose build, match, and output metrics.
contract that handles signed zero correctly.
batch boundaries, unmatched rows, invalid contracts, shared-buffer memory
accounting, multi-partition broadcast execution, and float-key rejection.
Are these changes tested?
Yes:
cargo fmt --allcargo clippy --all-targets --all-features -- -D warningscargo test -p datafusion-physical-plan joins::asof_join --all-featuresAre there any user-facing changes?
This adds a new physical operator API. The base operator deliberately rejects
floating-point equality keys; #24375 adds full Float16, Float32, and Float64
support. SQL and DataFrame APIs are left to later dependent PRs.