patrickswedish opened a new pull request, #26006:
URL: https://github.com/apache/datafusion/pull/26006
# 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_match`
verifies key equality by calling `compare_join_arrays`.
Prior to this change:
1. `compare_join_arrays` performed a manual downcast across a restricted
subset of primitive Arrow types and returned `not_impl_err!("Unsupported data
type in sort merge join comparator: {}", dt)` for other data types such as:
- Timezone-aware timestamps (`DataType::Timestamp(_, Some(_))`)
- Dictionary-encoded keys (`DataType::Dictionary(..)`)
- `DataType::Decimal256(..)`
- Time keys (`DataType::Time32(..)`, `DataType::Time64(..)`)
- Duration and Interval keys
2. As a consequence, queries performing semi, anti, or mark sort-merge joins
over timezone-aware timestamps or dictionary keys failed at runtime during
batch continuation with an "Unsupported data type" error.
3. In addition, `keys_match` previously contained an ad-hoc branch for
floating-point arrays to handle IEEE-754 `-0.0` normalization and NaN
comparisons via `JoinKeyComparator`, while other types fell back to
`compare_join_arrays`.
# What changes are included in this PR?
1. **Route `compare_join_arrays` directly to `JoinKeyComparator`**:
- Replaced the manual match table in `compare_join_arrays` with
`JoinKeyComparator::new(left_arrays, right_arrays, sort_options,
null_equality)?.compare(left, right)`.
- `JoinKeyComparator` builds Arrow's `make_comparator` closures once per
batch pair, providing support for all Arrow data types, preserving IEEE-754
`-0.0` normalization, and respecting `NullEquality` modes consistently across
all join streams.
2. **Clean up `keys_match` in `bitwise_stream.rs`**:
- Slices the first row of `right_arrays` and evaluates
`JoinKeyComparator::new(left_arrays, &right_keys, sort_options,
null_equality)?.is_equal(0, 0)`.
- Eliminates the special-cased float branch in `keys_match` since all
types are now handled uniformly.
3. **Tests**:
- Added unit test `test_compare_join_arrays_extended_types` in
`joins::utils::tests` validating comparisons across timezone-aware timestamps
(`TimestampMicrosecondArray` with `"UTC"`), dictionary arrays
(`Dictionary(Int32, Utf8)`), `Decimal256Array`, `Time64NanosecondArray`, and
`Time32MillisecondArray`.
# Are these changes tested?
Yes:
- New unit test `test_compare_join_arrays_extended_types` added in
`datafusion/physical-plan/src/joins/utils.rs`.
- All 247 existing sort-merge join tests pass (`cargo test --package
datafusion-physical-plan --lib joins::sort_merge_join::tests`).
- Formatted with `rustfmt` and verified with `cargo 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.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
