jayzhan211 commented on code in PR #26835:
URL: https://github.com/apache/datafusion/pull/26835#discussion_r4413009026
##########
datafusion/expr/src/udf.rs:
##########
@@ -1015,6 +1030,26 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send
+ Sync + Any {
false
}
+ /// Returns whether this function has been audited for range-key locality
+ /// analysis with these argument types and literal values.
+ ///
+ /// The safety requirements of
[`Self::supports_range_key_locality_analysis`]
+ /// apply to every admitted combination. Override this method when
admission
+ /// depends on literal values rather than argument types alone.
+ ///
+ /// `literal_arguments` has one entry per argument. `None` means the
argument
+ /// is not a literal; `Some(value)` supplies its value, which may itself be
+ /// NULL. Non-literal expressions are not evaluated to obtain these values.
+ ///
+ /// The default delegates to the type-only hook, preserving existing
opt-ins.
+ fn supports_range_key_locality_analysis_with_args(
Review Comment:
@benbellick thanks for this! The type-only
`supports_range_key_locality_analysis` hook from #26449 hasn't been released
yet (it's not on `branch-55`). So instead of adding a second hook that
delegates to it, could we add `literal_arguments` to that hook directly? That
leaves one method for implementors, removes both `_with_args` forwarding
wrappers (`ScalarUDF` and `AliasedScalarUDFImpl`), and needs only a signature
change in `date_trunc`. Once 56 ships, folding the two together would be a
breaking change, so it's easiest to do here. I tried it locally: about −21
lines, and the date_bin/date_trunc/partitioning/udf tests and clippy all pass.
```diff
- fn supports_range_key_locality_analysis(&self, _argument_types:
&[DataType]) -> bool {
- false
- }
-
- /// Returns whether this function has been audited for range-key
locality
- /// analysis with these argument types and literal values.
- ///
- /// The safety requirements of
[`Self::supports_range_key_locality_analysis`]
- /// apply to every admitted combination. Override this method when
admission
- /// depends on literal values rather than argument types alone.
- ///
- /// `literal_arguments` has one entry per argument. `None` means the
argument
- /// is not a literal; `Some(value)` supplies its value, which may
itself be
- /// NULL. Non-literal expressions are not evaluated to obtain these
values.
///
- /// The default delegates to the type-only hook, preserving existing
opt-ins.
- fn supports_range_key_locality_analysis_with_args(
+ /// `literal_arguments` has one entry per argument: `Some(value)` for a
+ /// literal (which may itself be NULL) and `None` otherwise.
+ fn supports_range_key_locality_analysis(
&self,
- argument_types: &[DataType],
+ _argument_types: &[DataType],
_literal_arguments: &[Option<&ScalarValue>],
) -> bool {
- self.supports_range_key_locality_analysis(argument_types)
+ false
}
```
and in `date_trunc.rs`:
```diff
- fn supports_range_key_locality_analysis(&self, argument_types:
&[DataType]) -> bool {
+ fn supports_range_key_locality_analysis(
+ &self,
+ argument_types: &[DataType],
+ _literal_arguments: &[Option<&ScalarValue>],
+ ) -> bool {
```
##########
datafusion/functions/src/datetime/date_bin.rs:
##########
@@ -289,6 +289,25 @@ impl ScalarUDFImpl for DateBinFunc {
}
}
+ fn supports_range_key_locality_analysis_with_args(
+ &self,
+ argument_types: &[DataType],
+ literal_arguments: &[Option<&ScalarValue>],
+ ) -> bool {
+ // With the epoch origin, a positive nanosecond stride rounds down
+ // monotonically. Timestamp binning preserves NULLs or returns an
error.
+ match (argument_types, literal_arguments) {
+ (
+ [
+ DataType::Interval(MonthDayNano),
+ Timestamp(Nanosecond, None),
+ ],
+ [Some(ScalarValue::IntervalMonthDayNano(Some(stride))), None],
+ ) => stride.months == 0 && stride.days == 0 && stride.nanoseconds
> 0,
Review Comment:
Optional: `date_bin_impl` converts `days` to a fixed 24 hours, so for
timezone-less input a day stride bins exactly like the same number of
nanoseconds. Admitting it would cover `INTERVAL '1 day'`, which is probably the
most common bucket. Computing the stride the same way `date_bin_impl` does
keeps the two in sync:
```rs
) => {
stride.months == 0
&& TimeDelta::try_days(i64::from(stride.days))
.and_then(|days|
days.checked_add(&Duration::nanoseconds(stride.nanoseconds)))
.and_then(|stride| stride.num_nanoseconds())
.is_some_and(|stride| stride > 0)
}
```
With these cases in
`range_key_locality_requires_positive_nanosecond_only_stride`, the tests pass
locally: `(0, 1, 0)` → true, `(0, 1, -1)` → true, `(0, -1, 1)` → false, `(1, 1,
0)` → false.
--
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]
