diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index a7332d2099d..4513d386f05 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -18,14 +18,13 @@ //! temporal kernels use chrono::{ - DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, TimeZone, Timelike, Utc, + DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, Offset, TimeZone, Timelike, }; use std::sync::Arc; use arrow::array::{ downcast_dictionary_array, downcast_temporal_array, - temporal_conversions::*, timezone::Tz, types::{ArrowDictionaryKeyType, ArrowTemporalType, TimestampMicrosecondType}, ArrowNumericType, @@ -35,6 +34,7 @@ use arrow::{ datatypes::{DataType, TimeUnit}, }; +use crate::utils::resolve_local_datetime; use crate::SparkError; // Copied from arrow_arith/temporal.rs @@ -101,54 +101,64 @@ fn trunc_days_to_week(days: i32) -> Option { Some(days - days_since_monday) } -// Based on arrow_arith/temporal.rs:extract_component_from_datetime_array -// Transforms an array of DateTime to an array of TimestampMicrosecond after applying an -// operation. The output array carries the input timezone annotation so downstream operators -// (shuffle, sort, row converter) observe a matching schema. -fn as_timestamp_tz_with_op, T: ArrowTemporalType, F>( - iter: ArrayIter, - mut builder: PrimitiveBuilder, - tz_str: &str, - op: F, -) -> Result -where - F: Fn(DateTime) -> i64, - i64: From, -{ - let tz: Tz = tz_str.parse()?; - for value in iter { - match value { - Some(value) => match as_datetime_with_timezone::(value.into(), tz) { - Some(time) => builder.append_value(op(time)), - _ => { - return Err(SparkError::Internal( - "Unable to read value as datetime".to_string(), - )); - } - }, - None => builder.append_null(), +/// How `date_trunc` truncates a timestamp with a timezone. Spark's `DateTimeUtils.truncTimestamp` +/// treats the levels differently, and matching it matters around DST transitions. +#[derive(Clone, Copy)] +enum TzTrunc { + /// `MICROSECOND`, `MILLISECOND` and `SECOND`. Offsets are whole seconds, so Spark truncates the + /// instant itself. The value is the unit in microseconds. + Instant(i64), + /// `MINUTE`, `HOUR` and `DAY`. Spark uses `ZonedDateTime.truncatedTo`, which truncates the local + /// time and keeps the input's offset if the result is ambiguous. + LocalTime(NtzTruncFn), + /// `WEEK`, `MONTH`, `QUARTER` and `YEAR`. Spark truncates the local date and then takes + /// `LocalDate.atStartOfDay`, which uses the earlier offset if midnight is ambiguous. + LocalDate(NtzTruncFn), +} + +/// Truncates `micros` in `tz` the way Spark's `DateTimeUtils.truncTimestamp` does. A truncated +/// local time that falls in a DST gap takes the offset from before the gap, which gives the same +/// instant as Java moving it forward by the gap's length. For the date levels that is also where +/// `atStartOfDay` puts a day whose midnight falls in a gap that starts at midnight. Returns `None` +/// if `micros` is out of chrono's range. +fn trunc_timestamp_in_tz(micros: i64, tz: &Tz, trunc: TzTrunc) -> Option { + let (trunc_fn, keep_offset) = match trunc { + TzTrunc::Instant(unit) => return Some(micros - micros.rem_euclid(unit)), + TzTrunc::LocalTime(trunc_fn) => (trunc_fn, true), + TzTrunc::LocalDate(trunc_fn) => (trunc_fn, false), + }; + let utc = DateTime::from_timestamp_micros(micros)?.naive_utc(); + let input_offset = tz.offset_from_utc_datetime(&utc).fix(); + let local = trunc_fn(utc.checked_add_offset(input_offset)?)?; + let truncated = match tz.offset_from_local_datetime(&local) { + LocalResult::Single(offset) => local.checked_sub_offset(offset.fix())?, + LocalResult::Ambiguous(earlier, later) => { + let offset = if keep_offset && later.fix() == input_offset { + later + } else { + earlier + }; + local.checked_sub_offset(offset.fix())? } - } - Ok(builder.finish().with_timezone(tz_str)) + LocalResult::None => resolve_local_datetime(tz, local).naive_utc(), + }; + Some(truncated.and_utc().timestamp_micros()) } -fn as_timestamp_tz_with_op_single( - value: Option, +/// Truncates one timezone-aware value and appends it to `builder`. +fn append_trunc_in_tz( builder: &mut PrimitiveBuilder, + value: Option, tz: &Tz, - op: F, -) -> Result<(), SparkError> -where - F: Fn(DateTime) -> i64, - i64: From, -{ + trunc: TzTrunc, +) -> Result<(), SparkError> { match value { - Some(value) => match as_datetime_with_timezone::(value.into(), *tz) { - Some(time) => builder.append_value(op(time)), - _ => { + Some(micros) => match trunc_timestamp_in_tz(micros, tz, trunc) { + Some(truncated) => builder.append_value(truncated), + None => { return Err(SparkError::Internal( "Unable to read value as datetime".to_string(), - )); + )) } }, None => builder.append_null(), @@ -156,25 +166,6 @@ where Ok(()) } -// Apply the Tz to the Naive Date Time, convert to UTC, and return as microseconds in Unix epoch. -// After truncation the carried UTC offset may be wrong if the truncated time falls in a different -// DST period than the original (e.g., truncating a December/PST timestamp to QUARTER yields -// October 1 which is in PDT). We re-resolve the naive local time through the timezone so that -// chrono picks the correct offset for the target date. -#[inline] -fn as_micros_from_unix_epoch_utc(dt: Option>) -> i64 { - let dt = dt.unwrap(); - let naive = dt.naive_local(); - let tz = dt.timezone(); - - match tz.from_local_datetime(&naive) { - LocalResult::Single(resolved) | LocalResult::Ambiguous(resolved, _) => { - resolved.with_timezone(&Utc).timestamp_micros() - } - LocalResult::None => dt.with_timezone(&Utc).timestamp_micros(), - } -} - #[inline] fn trunc_date_to_year(dt: T) -> Option { Some(dt) @@ -581,9 +572,6 @@ fn naive_to_micros(dt: NaiveDateTime) -> i64 { /// Truncates a `NaiveDateTime`, returning `None` if the result is out of range. type NtzTruncFn = fn(NaiveDateTime) -> Option; -/// Truncates a `DateTime`, returning `None` if the result is out of range. -type TzTruncFn = fn(DateTime) -> Option>; - /// The `timestamp_trunc` formats Spark accepts for the NTZ path, and the truncation each one /// selects. All entries are ASCII, so `eq_ignore_ascii_case` on the raw input matches Spark /// without allocating. @@ -605,23 +593,23 @@ const TIMESTAMP_TRUNC_FORMATS_NTZ: [(&str, NtzTruncFn); 15] = [ ("MICROSECOND", trunc_date_to_microsec), ]; -/// Same formats as `TIMESTAMP_TRUNC_FORMATS_NTZ`, monomorphized for the timezone-aware path. -const TIMESTAMP_TRUNC_FORMATS_TZ: [(&str, TzTruncFn); 15] = [ - ("YEAR", trunc_date_to_year), - ("YYYY", trunc_date_to_year), - ("YY", trunc_date_to_year), - ("QUARTER", trunc_date_to_quarter), - ("MONTH", trunc_date_to_month), - ("MON", trunc_date_to_month), - ("MM", trunc_date_to_month), - ("WEEK", trunc_date_to_week), - ("DAY", trunc_date_to_day), - ("DD", trunc_date_to_day), - ("HOUR", trunc_date_to_hour), - ("MINUTE", trunc_date_to_minute), - ("SECOND", trunc_date_to_second), - ("MILLISECOND", trunc_date_to_ms), - ("MICROSECOND", trunc_date_to_microsec), +/// The same formats as `TIMESTAMP_TRUNC_FORMATS_NTZ`, for the timezone-aware path. +const TIMESTAMP_TRUNC_FORMATS_TZ: [(&str, TzTrunc); 15] = [ + ("YEAR", TzTrunc::LocalDate(trunc_date_to_year)), + ("YYYY", TzTrunc::LocalDate(trunc_date_to_year)), + ("YY", TzTrunc::LocalDate(trunc_date_to_year)), + ("QUARTER", TzTrunc::LocalDate(trunc_date_to_quarter)), + ("MONTH", TzTrunc::LocalDate(trunc_date_to_month)), + ("MON", TzTrunc::LocalDate(trunc_date_to_month)), + ("MM", TzTrunc::LocalDate(trunc_date_to_month)), + ("WEEK", TzTrunc::LocalDate(trunc_date_to_week)), + ("DAY", TzTrunc::LocalTime(trunc_date_to_day)), + ("DD", TzTrunc::LocalTime(trunc_date_to_day)), + ("HOUR", TzTrunc::LocalTime(trunc_date_to_hour)), + ("MINUTE", TzTrunc::LocalTime(trunc_date_to_minute)), + ("SECOND", TzTrunc::Instant(1_000_000)), + ("MILLISECOND", TzTrunc::Instant(1_000)), + ("MICROSECOND", TzTrunc::Instant(1)), ]; /// Resolve a truncation format string to the corresponding NaiveDateTime truncation function. @@ -641,7 +629,7 @@ fn ntz_trunc_fn_for_format(format: &str) -> Result { } /// Timezone-aware sibling of `ntz_trunc_fn_for_format`. -fn tz_trunc_fn_for_format(format: &str) -> Result { +fn tz_trunc_for_format(format: &str) -> Result { TIMESTAMP_TRUNC_FORMATS_TZ .iter() .find(|(name, _)| name.eq_ignore_ascii_case(format)) @@ -711,18 +699,19 @@ where T: ArrowTemporalType + ArrowNumericType, i64: From, { - let builder = TimestampMicrosecondBuilder::with_capacity(array.len()); - let iter = ArrayIter::new(array); match array.data_type() { DataType::Timestamp(TimeUnit::Microsecond, None) => { // TimestampNTZ: operate directly on naive microsecond values without timezone timestamp_trunc_ntz(array, format) } - DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { - let trunc_fn = tz_trunc_fn_for_format(&format)?; - as_timestamp_tz_with_op::<&PrimitiveArray, T, _>(iter, builder, tz, |dt| { - as_micros_from_unix_epoch_utc(trunc_fn(dt)) - }) + DataType::Timestamp(TimeUnit::Microsecond, Some(tz_str)) => { + let trunc = tz_trunc_for_format(&format)?; + let tz: Tz = tz_str.parse()?; + let mut builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + for value in array.iter() { + append_trunc_in_tz(&mut builder, value.map(i64::from), &tz, trunc)?; + } + Ok(builder.finish().with_timezone(tz_str.as_ref())) } dt => return_compute_error_with!( "Unsupported input type '{:?}' for function 'timestamp_trunc'", @@ -826,10 +815,8 @@ macro_rules! timestamp_trunc_array_fmt_helper { DataType::Timestamp(TimeUnit::Microsecond, Some(tz_str)) => { let tz: Tz = tz_str.parse()?; for (index, val) in iter.enumerate() { - let trunc_fn = tz_trunc_fn_for_format($formats.value(index))?; - as_timestamp_tz_with_op_single::(val, &mut builder, &tz, |dt| { - as_micros_from_unix_epoch_utc(trunc_fn(dt)) - })?; + let trunc = tz_trunc_for_format($formats.value(index))?; + append_trunc_in_tz(&mut builder, val.map(|v| i64::from(v)), &tz, trunc)?; } Ok(builder.finish().with_timezone(tz_str.as_ref())) } @@ -1238,6 +1225,52 @@ mod tests { /// pre-fix kernel reused the input's MST offset for the truncated date, producing a result /// one hour late. Also verifies the output array carries the input timezone, which is what /// allows the result to flow through shuffle/sort without a `RowConverter` schema mismatch. + /// Truncation around DST transitions, against java.time, which Spark's `truncTimestamp` uses. + /// The ambiguous hours check that `MINUTE`, `HOUR` and `DAY` keep the input's offset, and São + /// Paulo on 2018-11-04 has no midnight. + #[test] + fn test_timestamp_trunc_matches_spark_at_dst_transitions() { + let formats = ["MINUTE", "HOUR", "DAY", "WEEK", "MONTH", "QUARTER", "YEAR"]; + #[rustfmt::skip] + let cases: [(&str, i64, [i64; 7]); 7] = [ + // 01:30 PDT on 2024-11-03, the first 01:30 of the fall-back day + ("America/Los_Angeles", 1_730_622_600_000_000, [1_730_622_600_000_000, 1_730_620_800_000_000, + 1_730_617_200_000_000, 1_730_098_800_000_000, 1_730_444_400_000_000, 1_727_766_000_000_000, + 1_704_096_000_000_000]), + // 01:30 PST on 2024-11-03, the second 01:30 + ("America/Los_Angeles", 1_730_626_200_000_000, [1_730_626_200_000_000, 1_730_624_400_000_000, + 1_730_617_200_000_000, 1_730_098_800_000_000, 1_730_444_400_000_000, 1_727_766_000_000_000, + 1_704_096_000_000_000]), + // 03:30 PDT on 2024-03-10, just after the spring-forward gap + ("America/Los_Angeles", 1_710_066_600_000_000, [1_710_066_600_000_000, 1_710_064_800_000_000, + 1_710_057_600_000_000, 1_709_539_200_000_000, 1_709_280_000_000_000, 1_704_096_000_000_000, + 1_704_096_000_000_000]), + // 01:39 PST on 1970-10-25, the second 01:39 of that fall-back day + ("America/Los_Angeles", 25_695_540_000_000, [25_695_540_000_000, 25_693_200_000_000, + 25_686_000_000_000, 25_167_600_000_000, 23_612_400_000_000, 23_612_400_000_000, + 28_800_000_000]), + // 13:00 on 2018-11-04 in São Paulo, a day whose midnight was skipped + ("America/Sao_Paulo", 1_541_343_600_000_000, [1_541_343_600_000_000, 1_541_343_600_000_000, + 1_541_300_400_000_000, 1_540_782_000_000_000, 1_541_041_200_000_000, 1_538_362_800_000_000, + 1_514_772_000_000_000]), + // 23:30 -02:00 on 2019-02-16, the first 23:30 of that fall-back night + ("America/Sao_Paulo", 1_550_367_000_000_000, [1_550_367_000_000_000, 1_550_365_200_000_000, + 1_550_282_400_000_000, 1_549_850_400_000_000, 1_548_986_400_000_000, 1_546_308_000_000_000, + 1_546_308_000_000_000]), + // 23:30 -03:00 on 2019-02-16, the second 23:30 + ("America/Sao_Paulo", 1_550_370_600_000_000, [1_550_370_600_000_000, 1_550_368_800_000_000, + 1_550_282_400_000_000, 1_549_850_400_000_000, 1_548_986_400_000_000, 1_546_308_000_000_000, + 1_546_308_000_000_000]), + ]; + for (tz, micros, expected) in cases { + let input = TimestampMicrosecondArray::from(vec![micros]).with_timezone(tz); + for (format, want) in formats.iter().zip(expected) { + let result = timestamp_trunc(&input, format.to_string()).unwrap(); + assert_eq!(result.value(0), want, "{tz} {micros} {format}"); + } + } + } + #[test] fn test_timestamp_trunc_dst_boundary() { // 2023-11-15 18:30:00 UTC = 2023-11-15 11:30 MST diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql index a1059736b9c..855122ac9a0 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_ambiguous.sql @@ -18,8 +18,10 @@ -- Regression test for ambiguous local times during DST fall-back. -- On 2024-11-03 at 2:00 AM America/Los_Angeles, clocks fall back to 1:00 AM, -- so 1:30 AM occurs twice (once in PDT, once in PST). Truncating 01:30 to HOUR --- gives 01:00, which is ambiguous. chrono's DateTime::with_minute(0) returns --- None for ambiguous results, causing a panic in as_micros_from_unix_epoch_utc. +-- gives 01:00, which is ambiguous. Spark's ZonedDateTime.truncatedTo keeps the +-- input's offset, so the PDT 01:30 truncates to 01:00 PDT and the PST 01:30 to +-- 01:00 PST. This used to panic natively: +-- https://github.com/apache/datafusion-comet/issues/5633 -- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true -- Config: spark.sql.session.timeZone=America/Los_Angeles @@ -29,22 +31,29 @@ CREATE TABLE test_trunc_ambiguous(ts timestamp) USING parquet statement INSERT INTO test_trunc_ambiguous VALUES - (timestamp('2024-11-03 01:30:00')) + (timestamp('2024-11-03 01:30:00')), + (timestamp('2024-11-03 01:30:00-07:00')), + (timestamp('2024-11-03 01:30:00-08:00')), + (timestamp('2024-11-03 01:00:00-08:00')), + (timestamp('1970-10-25 01:39:00-08:00')) -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) +query +SELECT ts, date_trunc('MINUTE', ts) FROM test_trunc_ambiguous ORDER BY ts + +query SELECT ts, date_trunc('DAY', ts) FROM test_trunc_ambiguous ORDER BY ts -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) +query SELECT ts, date_trunc('HOUR', ts) FROM test_trunc_ambiguous ORDER BY ts -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) +query SELECT ts, date_trunc('WEEK', ts) FROM test_trunc_ambiguous ORDER BY ts -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) +query SELECT ts, date_trunc('MONTH', ts) FROM test_trunc_ambiguous ORDER BY ts -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) +query SELECT ts, date_trunc('QUARTER', ts) FROM test_trunc_ambiguous ORDER BY ts -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) +query SELECT ts, date_trunc('YEAR', ts) FROM test_trunc_ambiguous ORDER BY ts diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_midnight.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_midnight.sql new file mode 100644 index 00000000000..6bd969ca23d --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_dst_midnight.sql @@ -0,0 +1,46 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +-- DST transitions at midnight. America/Sao_Paulo skipped midnight on 2018-11-04, +-- so DAY truncation and the day's start fall in the gap, and 23:00-24:00 on +-- 2019-02-16 happened twice, once at -02:00 and once at -03:00. +-- https://github.com/apache/datafusion-comet/issues/5633 + +-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Config: spark.sql.session.timeZone=America/Sao_Paulo + +statement +CREATE TABLE test_trunc_dst_midnight(ts timestamp) USING parquet + +statement +INSERT INTO test_trunc_dst_midnight VALUES + (timestamp('2018-11-03 23:30:00-03:00')), + (timestamp('2018-11-04 01:30:00-02:00')), + (timestamp('2018-11-04 13:00:00-02:00')), + (timestamp('2019-02-16 23:30:00-02:00')), + (timestamp('2019-02-16 23:30:00-03:00')), + (timestamp('2019-02-17 00:30:00-03:00')), + (NULL) + +query +SELECT ts, date_trunc('MINUTE', ts), date_trunc('HOUR', ts), date_trunc('DAY', ts) FROM test_trunc_dst_midnight ORDER BY ts + +query +SELECT ts, date_trunc('WEEK', ts), date_trunc('MONTH', ts) FROM test_trunc_dst_midnight ORDER BY ts + +query +SELECT ts, date_trunc('QUARTER', ts), date_trunc('YEAR', ts) FROM test_trunc_dst_midnight ORDER BY ts