jayzhan211 commented on code in PR #25635:
URL: https://github.com/apache/datafusion/pull/25635#discussion_r4176172834
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -435,10 +436,13 @@ impl ScalarUDFImpl for DateTruncFunc {
let precision = &input[0];
let date_value = &input[1];
- let order_safe_input = matches!(
- date_value.range.data_type(),
- Timestamp(_, None) | Time32(_) | Time64(_)
- );
+ let order_safe_input = match date_value.range.data_type() {
+ Timestamp(_, None) | Time32(_) | Time64(_) => true,
+ // UTC and fixed offsets have no transitions that could reverse
ordering.
+ Timestamp(_, Some(tz)) => is_fixed_offset_timezone(&tz),
Review Comment:
Dropping the sort makes `ORDER BY`/`GROUP BY
date_trunc('hour'|'day'|'month', ts)` over sorted UTC input 3–7x slower than
55.x. `SortPreservingMergeExec` (or `RepartitionExec preserve_order=true`) now
sits directly on the disjoint sorted file ranges and drains them one at a time.
That serializes the scan-embedded `date_trunc`, which runs chrono per row for
tz-aware hour and coarser (`time_elapsed_scanning_total` 10.87s vs 2.00s,
similar `elapsed_compute`).
Repro (release, 12 partitions):
```sql
COPY (SELECT arrow_cast(arrow_cast(value * 1000000000,
'Timestamp(Nanosecond, None)'),
'Timestamp(Nanosecond, Some("UTC"))') AS ts
FROM generate_series(0, 19999999) ORDER BY value) TO 'utc.parquet';
CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'utc.parquet' WITH ORDER
(ts ASC);
EXPLAIN ANALYZE SELECT date_trunc('hour', ts) AS d FROM t ORDER BY d;
-- main 0.51s, PR 1.56s
SELECT date_trunc('hour', ts) AS h, count(*) FROM t GROUP BY h ORDER BY h
LIMIT 1; -- main 0.20s, PR 1.39s
```
Fix: a fixed offset is UTC shifted by a constant, so truncate it
arithmetically (fine to land first as a separate PR). With this, hour/day/month
take 0.25/0.25/0.36s and GROUP BY hour 0.21s. Output matches the chrono path on
864k random values across all fixed offsets, units and hour–year granularities.
```diff
use chrono::{
- DateTime, Datelike, Duration, LocalResult, NaiveDateTime, Offset,
TimeDelta, Timelike,
+ DateTime, Datelike, Duration, LocalResult, NaiveDateTime, Offset,
TimeDelta,
+ TimeZone, Timelike,
};
@@ fn process_array<T: ArrowTimestampType>(
return Ok(ColumnarValue::Array(result));
}
+ // A fixed offset is UTC shifted by a constant: truncate local
time
+ // with timezone-free arithmetic, then shift back.
+ if let (Some(tz), Some(parsed)) = (tz_opt, parsed_tz)
+ && is_fixed_offset_timezone(tz)
+ {
+ let scale = match T::UNIT {
+ Second => 1,
+ Millisecond => 1_000,
+ Microsecond => 1_000_000,
+ Nanosecond => 1_000_000_000,
+ };
+ let offset = i64::from(
+ parsed
+
.offset_from_utc_datetime(&DateTime::UNIX_EPOCH.naive_utc())
+ .fix()
+ .local_minus_utc(),
+ ) * scale;
+ let out_of_range =
+ |x: i64| exec_datafusion_err!("Timestamp {x} out of
range");
+ let array: PrimitiveArray<T> = match granularity {
+ DatePart::Hour | DatePart::Day => {
+ let unit = scale
+ * if matches!(granularity, DatePart::Hour) {
+ 3_600
+ } else {
+ 86_400
+ };
+ array.try_unary(|x| {
+ x.checked_add(offset)
+ .and_then(|local|
x.checked_sub(local.rem_euclid(unit)))
+ .ok_or_else(|| out_of_range(x))
+ })?
+ }
+ _ => array.try_unary(|x| {
+ let local =
+ x.checked_add(offset).ok_or_else(||
out_of_range(x))?;
+ general_date_trunc(T::UNIT, local, None,
granularity)?
+ .checked_sub(offset)
+ .ok_or_else(|| out_of_range(x))
+ })?,
+ };
+ let array = array.with_timezone_opt(tz_opt.cloned());
+ return Ok(ColumnarValue::Array(Arc::new(array)));
+ }
+
let array: PrimitiveArray<T> = array
.try_unary(|x| general_date_trunc(T::UNIT, x, parsed_tz,
granularity))?
```
--
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]
