jayzhan211 commented on code in PR #25131:
URL: https://github.com/apache/datafusion/pull/25131#discussion_r3979409551
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -778,14 +778,15 @@ fn general_date_trunc_array_fine_granularity<T:
ArrowTimestampType>(
if let Some(unit) = unit {
let unit = unit.get();
- let array = PrimitiveArray::<T>::from_iter_values_with_nulls(
- array
- .values()
- .iter()
- .map(|v| *v - i64::rem_euclid(*v, unit)),
- array.nulls().cloned(),
- )
- .with_timezone_opt(tz_opt);
+ let array: PrimitiveArray<T> = array
+ .try_unary(|value| {
Review Comment:
**Perf: this de-vectorizes the fast path it's fixing.**
`try_unary` turns the straight-line map into a per-element `Result` with an
early return. Arrow says as much in the docs for the adjacent `try_unary_mut`:
*"LLVM is currently unable to effectively vectorize fallible operations"*. This
block is the fast path added in #23542 precisely to avoid per-value work, and
there's an in-repo bench for it:
```
cargo bench -p datafusion-functions --bench datetime_expressions --
date_trunc_minute_1000
```
PR vs `HEAD~1`: **706 ns → 829 ns, +16.4% (p = 0.00)**.
Underflow is only possible for values within one `unit` of `i64::MIN`, so it
can be detected with a branchless flag in the same pass, keeping the fallible
kernel for the rare array that trips it:
```diff
if let Some(unit) = unit {
let unit = unit.get();
- let array: PrimitiveArray<T> = array
- .try_unary(|value| {
- value.checked_sub(value.rem_euclid(unit)).ok_or_else(|| {
- exec_datafusion_err!(
- "Timestamp {value} out of range after truncating to
{granularity}"
- )
- })
- })?
- .with_timezone_opt(tz_opt);
+ // `value - value.rem_euclid(unit)` can only underflow within one
`unit`
+ // of `i64::MIN`. Detect that with a branchless flag so the common
case
+ // stays a single vectorizable pass, and only then take the
fallible path.
+ let bound = i64::MIN + unit;
+ let mut maybe_underflow = false;
+ let values: Vec<i64> = array
+ .values()
+ .iter()
+ .map(|v| {
+ maybe_underflow |= *v < bound;
+ v.wrapping_sub(v.rem_euclid(unit))
+ })
+ .collect();
+ let array: PrimitiveArray<T> = if !maybe_underflow {
+ PrimitiveArray::<T>::new(values.into(), array.nulls().cloned())
+ } else {
+ array.try_unary(|value| {
+ value.checked_sub(value.rem_euclid(unit)).ok_or_else(|| {
+ exec_datafusion_err!(
+ "Timestamp {value} out of range after truncating to
{granularity}"
+ )
+ })
+ })?
+ }
+ .with_timezone_opt(tz_opt);
Ok(Arc::new(array))
```
Measured **651 ns** on the same bench — 8.2% *faster* than main — with
`test_date_trunc_fine_granularity_underflow` and the rest of the `date_trunc`
tests passing, and `cargo clippy --all-targets --all-features` / `cargo fmt`
clean. (Null slots can hold arbitrary values, so a tripped fl; that's fine —
the fallback `try_unary` skips nulls and produces the right answer eitherway.)
--
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]
