namanjain24-sudo commented on code in PR #25326:
URL: https://github.com/apache/datafusion/pull/25326#discussion_r4052937018
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -778,26 +817,42 @@ fn general_date_trunc_array_fine_granularity<T:
ArrowTimestampType>(
if let Some(unit) = unit {
let unit = unit.get();
- // Truncation can only underflow within one `unit` of `i64::MIN`.
- // Track that possibility while computing the common case so the loop
- // remains infallible and can be vectorized.
- let underflow_bound = i64::MIN + unit;
+ // Truncate the value shifted by the (whole-unit) timezone offset and
+ // shift back: `truncate_tz(v) == truncate_naive(v + offset) - offset`.
+ // For fine granularities the offset is a multiple of `unit` and
cancels
+ // out; for hour/day it shifts the result as required.
+ let offset_in_unit = offset_nanos / nanos_per_unit(tu);
+ let add_lower = i64::MIN.saturating_sub(offset_in_unit);
+ let add_upper = i64::MAX.saturating_sub(offset_in_unit);
+ // Truncation of the shifted value can only underflow within one `unit`
+ // of `i64::MIN`. Track that (plus a wrapping shift at either extreme)
+ // while computing the common case so the loop remains infallible and
+ // can be vectorized.
+ let underflow_bound = i64::MIN + unit - offset_in_unit.min(0);
let mut maybe_underflow = false;
let values: Vec<i64> = array
.values()
.iter()
.map(|value| {
- maybe_underflow |= *value < underflow_bound;
- value.wrapping_sub(value.rem_euclid(unit))
+ maybe_underflow |=
+ *value < add_lower || *value > add_upper || *value <
underflow_bound;
+ let shifted = value.wrapping_add(offset_in_unit);
+ shifted.wrapping_sub(shifted.rem_euclid(unit)) - offset_in_unit
Review Comment:
This `- offset_in_unit` isn't wrapping, and it runs for every value before
`maybe_underflow` is checked, so it panics in debug builds near the ends of the
range. On this branch, `date_trunc('hour', ...)` over a `Timestamp(Second,
"+05:00")` column holding `i64::MIN` or `i64::MAX` panics here with `attempt to
subtract with overflow`; on `main` the same input returns `Timestamp ... out of
range`.
Since those values already set `maybe_underflow` and go through the checked
path, `.wrapping_sub(offset_in_unit)` should be enough. Could be worth adding
them to the underflow tests too.
--
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]
