sunchao commented on code in PR #6354:
URL: https://github.com/apache/datafusion-comet/pull/6354#discussion_r4127069931
##########
native/spark-expr/src/kernels/temporal.rs:
##########
@@ -101,80 +101,71 @@ fn trunc_days_to_week(days: i32) -> Option<i32> {
Some(days - days_since_monday)
}
-// Based on arrow_arith/temporal.rs:extract_component_from_datetime_array
-// Transforms an array of DateTime<Tz> 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<A: ArrayAccessor<Item = T::Native>, T:
ArrowTemporalType, F>(
- iter: ArrayIter<A>,
- mut builder: PrimitiveBuilder<TimestampMicrosecondType>,
- tz_str: &str,
- op: F,
-) -> Result<TimestampMicrosecondArray, SparkError>
-where
- F: Fn(DateTime<Tz>) -> i64,
- i64: From<T::Native>,
-{
- let tz: Tz = tz_str.parse()?;
- for value in iter {
- match value {
- Some(value) => match as_datetime_with_timezone::<T>(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<i64> {
+ 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(),
Review Comment:
[P2] Resolve `LocalDate` gaps at the first valid instant. With session
timezone `America/Toronto` and
`spark.comet.expression.TruncTimestamp.allowIncompatible=true`, apply
`date_trunc('WEEK', ts)` to a timestamp column containing `1919-03-31
00:45:00-04:00`. Spark returns `1919-03-31 00:30:00-04:00`, but this branch
returns `01:00:00-04:00`. Toronto's gap ran from the previous day's 23:30 to
00:30. `resolve_local_datetime` shifts midnight by the full gap length, whereas
Spark's `LocalDate.atStartOfDay` selects the gap's end. This replaces the
base's panic with silently incorrect weekly boundaries, even producing a result
later than this input. Keep gap-length shifting for `LocalTime`, but resolve
`LocalDate` to the transition boundary and add this regression case.
Evidence: Compiled the unmodified head/base kernels against locked Arrow
59.3.0 and chrono 0.4.45. For a timezone-annotated TimestampMicrosecondArray
containing -1601752500000000, timestamp_trunc(..., "WEEK") returns
-1601751600000000 at head. The base panics. Java 21's LocalDate.of(1919, 3,
31).atStartOfDay(ZoneId.of("America/Toronto")) returns -1601753400000000,
corresponding to 00:30-04:00. Spark's date-level truncation uses this
atStartOfDay rule through daysToMicros across the checked supported versions.
--
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]
