comphead commented on code in PR #6354: URL: https://github.com/apache/datafusion-comet/pull/6354#discussion_r4146853670
########## 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 Review Comment: Would it make sense to add a small `America/Havana` case where local midnight is ambiguous? Clocks fell back from 01:00 to 00:00 on 2020-11-01, and I don't think any current row makes `WEEK`, `MONTH`, `QUARTER` or `YEAR` land on a repeated local time, so the earlier-offset rule for those levels isn't pinned. I haven't run it, but for `timestamp('2020-11-01 00:30:00-05:00')` I'd expect `DAY` to give 00:00 at `-05:00` and `MONTH` to give 00:00 at `-04:00`. It would need its own session timezone, so probably a small new file. ########## 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(); Review Comment: Would it make sense to reuse `micros_to_naive` and `naive_to_micros` here? This line and the last line of the function look like the same conversions, and the NTZ path in this file already uses them. ########## native/spark-expr/src/kernels/temporal.rs: ########## @@ -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. Review Comment: Small thing: the new test seems to have landed between the existing `test_timestamp_trunc_dst_boundary` doc comment and its `#[test]` attribute. The Denver `QUARTER` lines above now read as part of this test's docs, and `test_timestamp_trunc_dst_boundary` has none. Would it make sense to move the new test above that comment block? -- 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]
