Repository navigation
fix: route sort-merge semi/anti/mark key equality through JoinKeyComparator - #26006
patrickswedish wants to merge 3 commits into
Conversation
…arator Sort-merge semi, anti, and mark joins (BitwiseSortMergeJoinStream) failed on timezone-aware timestamp, Dictionary, Decimal256, and Time keys with 'Unsupported data type in sort merge join comparator'. When a key group crosses batch boundaries, keys_match checks for key equality by calling compare_join_arrays. That function manually matched against a restricted subset of Arrow types and returned not_impl_err for everything else. Route compare_join_arrays and keys_match through JoinKeyComparator, which builds underlying Arrow make_comparator closures once per batch pair supporting all Arrow data types, preserving IEEE-754 -0.0 normalization and handling NullEquality modes consistently across all join streams. Closes apache#26003
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #26006 +/- ##
=======================================
Coverage 82.64% 82.64%
=======================================
Files 1147 1147
Lines 445502 445555 +53
Branches 445502 445555 +53
=======================================
+ Hits 368179 368231 +52
+ Misses 54969 54960 -9
- Partials 22354 22364 +10 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Thanks for this fix. Routing cross-batch key equality through JoinKeyComparator is appropriate and resolves the restricted type support in the previous comparison path. The changes are focused, and I found no correctness regression in the reviewed diff.
The architectural boundary is sound: JoinKeyComparator owns join-key comparison semantics, while the stream handles batch continuation. This keeps type support, NULL handling, and floating-point equality consistent with the merge scan without introducing another abstraction or duplicating comparison policy.
I'm approving because the implementation addresses the reported failure and the added tests cover representative extended types. I left a few non-blocking suggestions to strengthen execution-level and dictionary-encoding regression coverage and reduce test boilerplate.
| Ok(cmp == Ordering::Equal) | ||
| let right_keys = slice_keys(right_arrays, 0); | ||
| let cmp = | ||
| JoinKeyComparator::new(left_arrays, &right_keys, sort_options, null_equality)?; |
There was a problem hiding this comment.
Could we add a non-blocking execution-level regression using the existing join_float_key_batch_boundaries harness with timezone-aware timestamp keys? Please cover groups ending and continuing across batch boundaries, asserting semi/anti row IDs and mark bits, since the current tests exercise the helper rather than this stream's continuation path.
| ); | ||
|
|
||
| // 2. Dictionary(Int32, Utf8) | ||
| let dict_left: ArrayRef = Arc::new( |
There was a problem hiding this comment.
Could we also test equal dictionary values with different codes using independently ordered dictionaries built with DictionaryArray::try_new? The existing banana < cherry check already catches raw-code ordering, but the equal apple values both use code 0, so equality across different encodings remains untested. Please keep the existing inequality assertion.
| } | ||
|
|
||
| #[test] | ||
| fn test_compare_join_arrays_extended_types() -> Result<()> { |
There was a problem hiding this comment.
As a non-blocking cleanup, consider collecting the five (name, left, right) array pairs into a small table and running the shared equal/less assertions in a loop, with the case name in assertion messages. This would remove repeated comparator calls and make additional type coverage easier to maintain.
alamb
added
the
Prolific Author
|
This PR appears to be waiting on the author to respond to reviewer feedback so marking as draft. Please mark ready to review when it is ready for another look |
…egressions - Cover dictionary keys with independently ordered codes (different index encodings for equal values). - Consolidate extended type comparator test cases into a structured table loop. - Add execution-level join_timestamp_tz_key_batch_boundaries regression testing semi, anti, and mark joins across batch boundaries with continuing groups.
|
Thank you @kosiew for the thorough architectural review and approval! I have addressed all the non-blocking suggestions:
All tests pass locally. Marked ready for review @alamb. |
| } | ||
| } | ||
| Ok(res) | ||
| let cmp = |
There was a problem hiding this comment.
nit: nothing in this crate calls compare_join_arrays now, but it's still public. Building the comparator over the full arrays makes every call O(len): normalize_float_zero scans all the float values and logical_nulls scans dictionary keys. Slicing both sides to the one row first keeps it O(1) per call, like before:
- let cmp =
- JoinKeyComparator::new(left_arrays, right_arrays, sort_options, null_equality)?;
- Ok(cmp.compare(left, right))
+ let left_row: Vec<ArrayRef> = left_arrays.iter().map(|a| a.slice(left, 1)).collect();
+ let right_row: Vec<ArrayRef> =
+ right_arrays.iter().map(|a| a.slice(right, 1)).collect();
+ let cmp = JoinKeyComparator::new(&left_row, &right_row, sort_options, null_equality)?;
+ Ok(cmp.compare(0, 0))
…tor initialization Slice inputs to single-row array views prior to JoinKeyComparator construction so float normalization and dictionary key scanning remain constant-time rather than scanning full array buffers. Suggested-by: Jay Zhan <@jayzhan211>
|
Thank you @jayzhan211 for the thorough review, approval, and the sharp catch on You are completely right — constructing Applied your exact suggestion in |
There was a problem hiding this comment.
FWIW @patrickswedish sent my company a threat of litigation related to his contributions to this repository. While we think this has no merit, I would recommend we do not merge this PR or accept changes from this contributor, to avoid future headaches and legal claims
Which issue does this PR close?
Closes #26003.
Rationale for this change
In
BitwiseSortMergeJoinStream(which powers sort-merge semi, anti, and mark joins), when a key group spans across batch boundaries,keys_matchverifies key equality by callingcompare_join_arrays.Prior to this change:
compare_join_arraysperformed a manual downcast across a restricted subset of primitive Arrow types and returnednot_impl_err!("Unsupported data type in sort merge join comparator: {}", dt)for other data types such as:DataType::Timestamp(_, Some(_)))DataType::Dictionary(..))DataType::Decimal256(..)DataType::Time32(..),DataType::Time64(..))keys_matchpreviously contained an ad-hoc branch for floating-point arrays to handle IEEE-754-0.0normalization and NaN comparisons viaJoinKeyComparator, while other types fell back tocompare_join_arrays.What changes are included in this PR?
compare_join_arraysdirectly toJoinKeyComparator:compare_join_arrayswithJoinKeyComparator::new(left_arrays, right_arrays, sort_options, null_equality)?.compare(left, right).JoinKeyComparatorbuilds Arrow'smake_comparatorclosures once per batch pair, providing support for all Arrow data types, preserving IEEE-754-0.0normalization, and respectingNullEqualitymodes consistently across all join streams.keys_matchinbitwise_stream.rs:right_arraysand evaluatesJoinKeyComparator::new(left_arrays, &right_keys, sort_options, null_equality)?.is_equal(0, 0).keys_matchsince all types are now handled uniformly.test_compare_join_arrays_extended_typesinjoins::utils::testsvalidating comparisons across timezone-aware timestamps (TimestampMicrosecondArraywith"UTC"), dictionary arrays (Dictionary(Int32, Utf8)),Decimal256Array,Time64NanosecondArray, andTime32MillisecondArray.Are these changes tested?
Yes:
test_compare_join_arrays_extended_typesadded indatafusion/physical-plan/src/joins/utils.rs.cargo test --package datafusion-physical-plan --lib joins::sort_merge_join::tests).rustfmtand verified withcargo clippy --package datafusion-physical-plan --lib -- -D warnings(0 warnings).Are there any user-facing changes?
No breaking changes. Users can now execute sort-merge semi, anti, and mark joins over timezone-aware timestamps, dictionary keys,
Decimal256, and time keys without errors.