From 029dafd3e025930720494b05d6cef34f1f6600a4 Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Tue, 15 Sep 2026 18:39:15 +0800 Subject: [PATCH 1/3] feat: Use DataFusion date_trunc for compatible trunc and date_trunc paths Reuse the upstream physical date_trunc for scalar formats while keeping Comet per-row format, dictionary, and session-timezone handling, with a fallback outside DataFusion's TimestampNanosecond range. --- native/Cargo.lock | 1 + native/Cargo.toml | 1 + native/spark-expr/Cargo.toml | 1 + native/spark-expr/src/kernels/temporal.rs | 752 ++++++++++++++++-- .../expressions/datetime/trunc_date.sql | 49 +- .../expressions/datetime/trunc_timestamp.sql | 54 +- .../trunc_timestamp_dst_ambiguous.sql | 63 +- .../comet/CometTemporalExpressionSuite.scala | 46 +- 8 files changed, 852 insertions(+), 115 deletions(-) diff --git a/native/Cargo.lock b/native/Cargo.lock index fabb56e2ec8..9165ad24edc 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2100,6 +2100,7 @@ dependencies = [ "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", + "datafusion-functions", "datafusion-spark", "futures", "jni 0.22.4", diff --git a/native/Cargo.toml b/native/Cargo.toml index bde2ba432c0..1351e6b9ae4 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -44,6 +44,7 @@ bytes = { version = "1.11.1" } parquet = { version = "59.2.0", default-features = false, features = ["experimental"] } datafusion = { version = "55.1.0", default-features = false, features = ["unicode_expressions", "crypto_expressions", "nested_expressions", "parquet"] } datafusion-datasource = { version = "55.1.0" } +datafusion-functions = { version = "55.1.0", default-features = false, features = ["datetime_expressions"] } datafusion-physical-expr-adapter = { version = "55.1.0" } datafusion-spark = { version = "55.1.0", features = ["core"] } datafusion-comet-spark-expr = { path = "spark-expr" } diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 9019dffbb50..44829c8a85c 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -30,6 +30,7 @@ edition = { workspace = true } arrow = { workspace = true } chrono = { workspace = true } datafusion = { workspace = true } +datafusion-functions = { workspace = true } datafusion-spark = { workspace = true } num = { workspace = true } regex = { workspace = true } diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index a7332d2099d..b0f5fb200aa 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -32,11 +32,55 @@ use arrow::array::{ }; use arrow::{ array::*, - datatypes::{DataType, TimeUnit}, + compute::cast, + datatypes::{DataType, Field, TimeUnit}, }; +use datafusion::{ + common::{config::ConfigOptions, ScalarValue}, + logical_expr::{ColumnarValue, ScalarFunctionArgs}, +}; +use datafusion_functions::datetime; use crate::SparkError; +/// Invoke DataFusion's physical `date_trunc` implementation with a scalar granularity. +/// +/// Spark syntax normalization and compatibility fallback decisions deliberately live outside +/// this helper so the upstream execution boundary stays obvious. +fn datafusion_date_trunc( + array: ArrayRef, + granularity: &'static str, +) -> Result { + let data_type = array.data_type().clone(); + let number_rows = array.len(); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some(granularity.to_string()))), + ColumnarValue::Array(array), + ]; + let arg_fields = args + .iter() + .enumerate() + .map(|(index, value)| { + Arc::new(Field::new( + format!("date_trunc_arg_{index}"), + value.data_type(), + true, + )) + }) + .collect(); + + datetime::date_trunc() + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Arc::new(Field::new("date_trunc", data_type, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .and_then(|value| value.to_array(number_rows)) + .map_err(|error| SparkError::Internal(error.to_string())) +} + // Copied from arrow_arith/temporal.rs macro_rules! return_compute_error_with { ($msg:expr, $param:expr) => { @@ -324,6 +368,18 @@ where /// Truncates a date expressed as days since the epoch, returning `None` if it is out of range. type DateTruncFn = fn(i32) -> Option; +/// The Spark `trunc` spellings and the canonical granularities accepted by DataFusion. +const DATE_TRUNC_ALIASES: [(&str, &str); 8] = [ + ("YEAR", "year"), + ("YYYY", "year"), + ("YY", "year"), + ("QUARTER", "quarter"), + ("MONTH", "month"), + ("MON", "month"), + ("MM", "month"), + ("WEEK", "week"), +]; + /// The `date_trunc` formats Spark accepts, and the truncation each one selects. const DATE_TRUNC_FORMATS: [(&str, DateTruncFn); 8] = [ ("YEAR", trunc_days_to_year), @@ -353,19 +409,110 @@ fn date_trunc_fn_for_format(format: &str) -> Result { }) } -/// Optimized date truncation for Date32 arrays -/// Works directly with days since epoch instead of converting to/from NaiveDateTime +/// Normalize a Spark `trunc` format without exposing additional DataFusion granularities. +fn normalize_date_trunc_format(format: &str) -> Result<&'static str, SparkError> { + DATE_TRUNC_ALIASES + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case(format)) + .map(|(_, granularity)| *granularity) + .ok_or_else(|| { + SparkError::Internal(format!( + "Unsupported format: {format:?} for function 'date_trunc'" + )) + }) +} + +const MICROS_PER_DAY: i64 = 86_400_000_000; + +#[inline] +fn fits_timestamp_nanosecond(micros: i64) -> bool { + micros.checked_mul(1_000).is_some() +} + +/// DataFusion's coarse truncation first converts the input to nanoseconds. Although an input near +/// the lower TimestampNanosecond bound can itself be represented, truncating it may move the +/// result before that bound: for example, `1677-09-22` truncated to YEAR becomes `1677-01-01`. +/// The result can move backward by 365 days (366 when starting from December 31 in a leap year), +/// and timezone gap handling can shift it by a few more hours. Round that worst case up to 370 +/// days so both DataFusion's input and coarse-truncation result remain representable. The +/// effective microsecond interval is approximately `1678-09-26T00:12:43.145225Z` through +/// `2262-04-11T23:47:16.854775Z`; because Date32 values are UTC midnight, its first upstream date +/// is 1678-09-27 and its last is 2262-04-11. +#[inline] +fn fits_datafusion_coarse_trunc_range(micros: i64) -> bool { + const LOWER_NANOSECOND_MICROS: i64 = i64::MIN / 1_000; + const COARSE_TRUNC_MARGIN_MICROS: i64 = 370 * MICROS_PER_DAY; + + fits_timestamp_nanosecond(micros) + && micros >= LOWER_NANOSECOND_MICROS + COARSE_TRUNC_MARGIN_MICROS +} + +#[inline] +fn date32_to_utc_midnight_micros(days: i32) -> Option { + i64::from(days).checked_mul(MICROS_PER_DAY) +} + +#[inline] +fn date32_fits_upstream(days: i32) -> bool { + date32_to_utc_midnight_micros(days).is_some_and(fits_datafusion_coarse_trunc_range) +} + +/// Truncate scalar-format Date32 values through DataFusion's physical `date_trunc`. +/// +/// DataFusion 55.1 scales coarse timestamp granularities to nanoseconds internally. Spark Date +/// supports approximately years 0001 through 9999, while TimestampNanosecond only spans roughly +/// 1677 through 2262. Values outside the guarded TimestampNanosecond range therefore retain the +/// established Date32 calculation; values inside it use the upstream cast sandwich. fn date_trunc_date32(array: &Date32Array, format: String) -> Result { - // Select the truncation function based on format + let granularity = normalize_date_trunc_format(&format)?; let trunc_fn = date_trunc_fn_for_format(&format)?; - - // Apply truncation to each element - let result: Date32Array = array + let mut has_wide_value = false; + let upstream_input: Date32Array = array .iter() - .map(|opt_days| opt_days.and_then(trunc_fn)) + .map(|value| { + value.and_then(|days| { + if date32_fits_upstream(days) { + Some(days) + } else { + has_wide_value = true; + None + } + }) + }) .collect(); - Ok(result) + if upstream_input.null_count() == array.len() { + return Ok(array.iter().map(|value| value.and_then(trunc_fn)).collect()); + } + + let timestamps = cast( + &upstream_input, + &DataType::Timestamp(TimeUnit::Microsecond, None), + )?; + let truncated = datafusion_date_trunc(timestamps, granularity)?; + let truncated = cast(truncated.as_ref(), &DataType::Date32)?; + let upstream = truncated + .as_any() + .downcast_ref::() + .expect("DataFusion date_trunc Date32 cast mismatch"); + + if !has_wide_value { + return Ok(upstream.clone()); + } + + Ok(array + .iter() + .enumerate() + .map(|(index, value)| { + value.and_then(|days| { + if date32_fits_upstream(days) { + Some(upstream.value(index)) + } else { + trunc_fn(days) + } + }) + }) + .collect()) } /// @@ -584,6 +731,25 @@ type NtzTruncFn = fn(NaiveDateTime) -> Option; /// Truncates a `DateTime`, returning `None` if the result is out of range. type TzTruncFn = fn(DateTime) -> Option>; +/// The Spark `date_trunc` spellings and their canonical DataFusion granularities. +const TIMESTAMP_TRUNC_ALIASES: [(&str, &str); 15] = [ + ("YEAR", "year"), + ("YYYY", "year"), + ("YY", "year"), + ("QUARTER", "quarter"), + ("MONTH", "month"), + ("MON", "month"), + ("MM", "month"), + ("WEEK", "week"), + ("DAY", "day"), + ("DD", "day"), + ("HOUR", "hour"), + ("MINUTE", "minute"), + ("SECOND", "second"), + ("MILLISECOND", "millisecond"), + ("MICROSECOND", "microsecond"), +]; + /// 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. @@ -653,6 +819,19 @@ fn tz_trunc_fn_for_format(format: &str) -> Result { }) } +/// Normalize Spark `date_trunc` aliases without accepting additional DataFusion spellings. +fn normalize_timestamp_trunc_format(format: &str) -> Result<&'static str, SparkError> { + TIMESTAMP_TRUNC_ALIASES + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case(format)) + .map(|(_, granularity)| *granularity) + .ok_or_else(|| { + SparkError::Internal(format!( + "Unsupported format: {format:?} for function 'timestamp_trunc'" + )) + }) +} + /// Truncate a TimestampNTZ array without any timezone conversion. /// NTZ values are timezone-independent; we treat the raw microseconds as a naive datetime. fn timestamp_trunc_ntz( @@ -680,6 +859,99 @@ where Ok(result) } +/// The scalar-format implementation retained for values outside DataFusion 55.1's internal +/// TimestampNanosecond range. Row-format paths continue to call the same underlying helpers. +fn timestamp_trunc_legacy( + array: &TimestampMicrosecondArray, + format: &str, +) -> Result { + let builder = TimestampMicrosecondBuilder::with_capacity(array.len()); + let iter = ArrayIter::new(array); + match array.data_type() { + DataType::Timestamp(TimeUnit::Microsecond, None) => { + timestamp_trunc_ntz(array, format.to_string()) + } + DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => { + let trunc_fn = tz_trunc_fn_for_format(format)?; + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + iter, + builder, + tz, + |dt| as_micros_from_unix_epoch_utc(trunc_fn(dt)), + ) + } + dt => return_compute_error_with!( + "Unsupported input type '{:?}' for function 'timestamp_trunc'", + dt + ), + } +} + +fn datafusion_timestamp_trunc_requires_nanos(granularity: &str, has_timezone: bool) -> bool { + match granularity { + "microsecond" | "millisecond" | "second" | "minute" => false, + "hour" | "day" => has_timezone, + "week" | "month" | "quarter" | "year" => true, + _ => unreachable!("granularity was normalized before compatibility dispatch"), + } +} + +fn timestamp_trunc_upstream( + array: &TimestampMicrosecondArray, + format: &str, +) -> Result { + let granularity = normalize_timestamp_trunc_format(format)?; + let requires_nanos = + datafusion_timestamp_trunc_requires_nanos(granularity, array.timezone().is_some()); + + if !requires_nanos + || array + .iter() + .flatten() + .all(fits_datafusion_coarse_trunc_range) + { + let result = datafusion_date_trunc(Arc::new(array.clone()), granularity)?; + return Ok(result + .as_any() + .downcast_ref::() + .expect("DataFusion date_trunc timestamp result mismatch") + .clone()); + } + + let upstream_input = TimestampMicrosecondArray::from_iter( + array + .iter() + .map(|value| value.filter(|micros| fits_datafusion_coarse_trunc_range(*micros))), + ) + .with_timezone_opt(array.timezone()); + let legacy_input = TimestampMicrosecondArray::from_iter( + array + .iter() + .map(|value| value.filter(|micros| !fits_datafusion_coarse_trunc_range(*micros))), + ) + .with_timezone_opt(array.timezone()); + + let upstream = datafusion_date_trunc(Arc::new(upstream_input), granularity)?; + let upstream = upstream + .as_any() + .downcast_ref::() + .expect("DataFusion date_trunc timestamp result mismatch"); + let legacy = timestamp_trunc_legacy(&legacy_input, format)?; + + Ok( + TimestampMicrosecondArray::from_iter(array.iter().enumerate().map(|(index, value)| { + value.map(|micros| { + if fits_datafusion_coarse_trunc_range(micros) { + upstream.value(index) + } else { + legacy.value(index) + } + }) + })) + .with_timezone_opt(array.timezone()), + ) +} + /// Truncate a single NTZ value and append to builder fn timestamp_trunc_ntz_single( value: Option, @@ -711,19 +983,14 @@ 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, _) => timestamp_trunc_upstream( + array + .as_any() + .downcast_ref::() + .expect("TimestampMicrosecond type mismatch"), + &format, + ), dt => return_compute_error_with!( "Unsupported input type '{:?}' for function 'timestamp_trunc'", dt @@ -897,37 +1164,248 @@ where #[cfg(test)] mod tests { use crate::kernels::temporal::{ - date_trunc, date_trunc_array_fmt_dyn, timestamp_trunc, timestamp_trunc_array_fmt_dyn, + date_trunc, date_trunc_array_fmt_dyn, date_trunc_dyn, timestamp_trunc, + timestamp_trunc_array_fmt_dyn, timestamp_trunc_dyn, }; + use crate::SparkError; use arrow::array::{ builder::{PrimitiveDictionaryBuilder, StringDictionaryBuilder}, iterator::ArrayIter, types::{Date32Type, Int32Type, TimestampMicrosecondType}, Array, Date32Array, PrimitiveArray, StringArray, TimestampMicrosecondArray, }; + use chrono::{DateTime, Datelike, NaiveDate}; use std::sync::Arc; + fn epoch_days(date: &str) -> i32 { + NaiveDate::parse_from_str(date, "%Y-%m-%d") + .unwrap() + .num_days_from_ce() + - 719_163 + } + + fn assert_date_trunc(format: &str, input: &[Option<&str>], expected: &[Option<&str>]) { + let input = Date32Array::from( + input + .iter() + .map(|date| date.map(epoch_days)) + .collect::>(), + ); + let expected = Date32Array::from( + expected + .iter() + .map(|date| date.map(epoch_days)) + .collect::>(), + ); + assert_eq!(date_trunc(&input, format.to_string()).unwrap(), expected); + } + + fn instant_micros(instant: &str) -> i64 { + DateTime::parse_from_rfc3339(instant) + .unwrap() + .timestamp_micros() + } + + fn assert_timestamp_trunc( + format: &str, + timezone: Option<&str>, + input: &[Option<&str>], + expected: &[Option<&str>], + ) { + let input = TimestampMicrosecondArray::from( + input + .iter() + .map(|instant| instant.map(instant_micros)) + .collect::>(), + ) + .with_timezone_opt(timezone); + let expected = TimestampMicrosecondArray::from( + expected + .iter() + .map(|instant| instant.map(instant_micros)) + .collect::>(), + ) + .with_timezone_opt(timezone); + assert_eq!( + timestamp_trunc(&input, format.to_string()).unwrap(), + expected + ); + } + #[test] - #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_date_trunc() { - let size = 1000; - let mut vec: Vec = Vec::with_capacity(size); - for i in 0..size { - vec.push(i as i32); + for format in ["YEAR", "YYYY", "YY", "year", "Year", "yEaR"] { + assert_date_trunc(format, &[Some("2024-05-17")], &[Some("2024-01-01")]); } - let array = Date32Array::from(vec); - for fmt in [ - "YEAR", "YYYY", "YY", "QUARTER", "MONTH", "MON", "MM", "WEEK", - ] { - match date_trunc(&array, fmt.to_string()) { - Ok(a) => { - for i in 0..size { - assert!(array.values().get(i) >= a.values().get(i)) - } - } - _ => unreachable!(), - } + for format in ["MONTH", "MON", "MM", "month", "Mon"] { + assert_date_trunc(format, &[Some("2024-05-17")], &[Some("2024-05-01")]); } + + assert_date_trunc( + "QUARTER", + &[ + Some("2024-01-01"), + Some("2024-03-31"), + Some("2024-04-01"), + Some("2024-06-30"), + Some("2024-07-01"), + Some("2024-10-01"), + ], + &[ + Some("2024-01-01"), + Some("2024-01-01"), + Some("2024-04-01"), + Some("2024-04-01"), + Some("2024-07-01"), + Some("2024-10-01"), + ], + ); + + assert_date_trunc( + "week", + &[ + Some("2024-05-13"), + Some("2024-05-14"), + Some("2024-05-19"), + Some("2024-05-01"), + Some("2024-01-01"), + Some("2023-12-31"), + ], + &[ + Some("2024-05-13"), + Some("2024-05-13"), + Some("2024-05-13"), + Some("2024-04-29"), + Some("2024-01-01"), + Some("2023-12-25"), + ], + ); + + let dates = [ + Some("2024-02-29"), + Some("2000-02-29"), + Some("1900-02-28"), + Some("1969-12-31"), + Some("1960-02-29"), + Some("1900-01-01"), + // The input fits TimestampNanosecond, but truncating it to YEAR does not. + Some("1677-09-22"), + // Valid Spark Date32 outside TimestampNanosecond's range. DataFusion 55.1 + // date_trunc internally converts coarse granularities to nanoseconds. + Some("3333-05-17"), + None, + ]; + assert_date_trunc( + "YEAR", + &dates, + &[ + Some("2024-01-01"), + Some("2000-01-01"), + Some("1900-01-01"), + Some("1969-01-01"), + Some("1960-01-01"), + Some("1900-01-01"), + Some("1677-01-01"), + Some("3333-01-01"), + None, + ], + ); + assert_date_trunc( + "QUARTER", + &dates, + &[ + Some("2024-01-01"), + Some("2000-01-01"), + Some("1900-01-01"), + Some("1969-10-01"), + Some("1960-01-01"), + Some("1900-01-01"), + Some("1677-07-01"), + Some("3333-04-01"), + None, + ], + ); + assert_date_trunc( + "MONTH", + &dates, + &[ + Some("2024-02-01"), + Some("2000-02-01"), + Some("1900-02-01"), + Some("1969-12-01"), + Some("1960-02-01"), + Some("1900-01-01"), + Some("1677-09-01"), + Some("3333-05-01"), + None, + ], + ); + assert_date_trunc( + "WEEK", + &dates, + &[ + Some("2024-02-26"), + Some("2000-02-28"), + Some("1900-02-26"), + Some("1969-12-29"), + Some("1960-02-29"), + Some("1900-01-01"), + Some("1677-09-20"), + Some("3333-05-11"), + None, + ], + ); + + let input = Date32Array::from(vec![epoch_days("2024-05-17")]); + for format in ["DAY", "HOUR", "SECOND", "invalid", " YEAR ", ""] { + let SparkError::Internal(message) = date_trunc(&input, format.to_string()).unwrap_err() + else { + panic!("expected an internal unsupported-format error"); + }; + assert_eq!( + message, + format!("Unsupported format: {format:?} for function 'date_trunc'") + ); + } + + assert!(super::date32_to_utc_midnight_micros(i32::MAX).is_none()); + assert!(super::date32_to_utc_midnight_micros(i32::MIN).is_none()); + assert!(!super::date32_fits_upstream(epoch_days("3333-05-17"))); + assert!(!super::date32_fits_upstream(epoch_days("0001-01-01"))); + assert!(super::date32_fits_upstream(epoch_days("2024-05-17"))); + assert!(super::date32_fits_upstream(epoch_days("1678-09-27"))); + assert!(!super::date32_fits_upstream(epoch_days("1678-09-26"))); + } + + #[test] + fn test_date_trunc_scalar_format_dictionary() { + let mut builder = PrimitiveDictionaryBuilder::::new(); + builder.append(epoch_days("2024-05-17")).unwrap(); + builder.append(epoch_days("2024-06-30")).unwrap(); + builder.append(epoch_days("2024-05-17")).unwrap(); + builder.append_null(); + builder.append(epoch_days("1969-12-31")).unwrap(); + let input = builder.finish(); + let input_keys = input.keys().clone(); + + let result = date_trunc_dyn(&input, "MONTH".to_string()).unwrap(); + let result = result + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(result.keys(), &input_keys); + + let decoded = result.downcast_dict::().unwrap(); + assert_eq!( + decoded.into_iter().collect::>(), + vec![ + Some(epoch_days("2024-05-01")), + Some(epoch_days("2024-06-01")), + Some(epoch_days("2024-05-01")), + None, + Some(epoch_days("1969-12-01")), + ] + ); } #[test] @@ -1047,40 +1525,180 @@ mod tests { } #[test] - #[cfg_attr(miri, ignore)] // test takes too long with miri fn test_timestamp_trunc() { - let size = 1000; - let mut vec: Vec = Vec::with_capacity(size); - for i in 0..size { - vec.push(i as i64); + let input = [Some("2024-05-17T12:34:56.123456Z"), None]; + for format in ["YEAR", "YYYY", "YY", "year", "Year", "yEaR"] { + assert_timestamp_trunc( + format, + Some("UTC"), + &input, + &[Some("2024-01-01T00:00:00Z"), None], + ); } - let array = TimestampMicrosecondArray::from(vec).with_timezone_utc(); - for fmt in [ + for format in ["MONTH", "MON", "MM", "month", "Mon"] { + assert_timestamp_trunc( + format, + Some("UTC"), + &input, + &[Some("2024-05-01T00:00:00Z"), None], + ); + } + for (format, expected) in [ + ("QUARTER", "2024-04-01T00:00:00Z"), + ("WEEK", "2024-05-13T00:00:00Z"), + ("DAY", "2024-05-17T00:00:00Z"), + ("DD", "2024-05-17T00:00:00Z"), + ("HOUR", "2024-05-17T12:00:00Z"), + ("MINUTE", "2024-05-17T12:34:00Z"), + ("SECOND", "2024-05-17T12:34:56Z"), + ("MILLISECOND", "2024-05-17T12:34:56.123Z"), + ("MICROSECOND", "2024-05-17T12:34:56.123456Z"), + ] { + assert_timestamp_trunc(format, Some("UTC"), &input, &[Some(expected), None]); + } + + let invalid_input = + TimestampMicrosecondArray::from(vec![instant_micros("2024-05-17T12:34:56Z")]) + .with_timezone_utc(); + for format in ["MILLISECONDS", "invalid", " DAY ", ""] { + let SparkError::Internal(message) = + timestamp_trunc(&invalid_input, format.to_string()).unwrap_err() + else { + panic!("expected an internal unsupported-format error"); + }; + assert_eq!( + message, + format!("Unsupported format: {format:?} for function 'timestamp_trunc'") + ); + } + } + + #[test] + fn test_timestamp_trunc_wide_range_fallback() { + let input = [ + Some("2024-05-17T12:34:56.123456Z"), + Some("3333-05-17T12:34:56.123456Z"), + Some("1969-12-31T23:59:59.123456Z"), + Some("1677-09-22T00:00:00Z"), + None, + ]; + assert_timestamp_trunc( "YEAR", - "YYYY", - "YY", + Some("UTC"), + &input, + &[ + Some("2024-01-01T00:00:00Z"), + Some("3333-01-01T00:00:00Z"), + Some("1969-01-01T00:00:00Z"), + Some("1677-01-01T00:00:00Z"), + None, + ], + ); + assert_timestamp_trunc( "QUARTER", - "MONTH", - "MON", - "MM", + Some("UTC"), + &input, + &[ + Some("2024-04-01T00:00:00Z"), + Some("3333-04-01T00:00:00Z"), + Some("1969-10-01T00:00:00Z"), + Some("1677-07-01T00:00:00Z"), + None, + ], + ); + assert_timestamp_trunc( "WEEK", - "DAY", - "DD", + Some("UTC"), + &input, + &[ + Some("2024-05-13T00:00:00Z"), + Some("3333-05-11T00:00:00Z"), + Some("1969-12-29T00:00:00Z"), + Some("1677-09-20T00:00:00Z"), + None, + ], + ); + } + + #[test] + fn test_timestamp_trunc_dst_gap_and_overlap() { + assert_timestamp_trunc( "HOUR", - "MINUTE", - "SECOND", - "MILLISECOND", - "MICROSECOND", - ] { - match timestamp_trunc(&array, fmt.to_string()) { - Ok(a) => { - for i in 0..size { - assert!(array.values().get(i) >= a.values().get(i)) - } - } - _ => unreachable!(), - } - } + Some("America/Los_Angeles"), + &[ + Some("2024-11-03T08:30:15.123456Z"), + Some("2024-11-03T09:30:15.123456Z"), + ], + &[Some("2024-11-03T08:00:00Z"), Some("2024-11-03T09:00:00Z")], + ); + assert_timestamp_trunc( + "DAY", + Some("America/New_York"), + &[ + Some("2024-03-10T06:30:15.123456Z"), + Some("2024-03-10T07:30:15.123456Z"), + Some("2024-11-03T06:30:15.123456Z"), + ], + &[ + Some("2024-03-10T05:00:00Z"), + Some("2024-03-10T05:00:00Z"), + Some("2024-11-03T04:00:00Z"), + ], + ); + // Sao Paulo advanced from 23:59:59 on November 3 directly to 01:00 on November 4. + // Spark resolves DAY for a post-gap instant to that day's first valid local time. + assert_timestamp_trunc( + "DAY", + Some("America/Sao_Paulo"), + &[ + Some("2018-11-04T02:30:15.123456Z"), + Some("2018-11-04T03:30:15.123456Z"), + ], + &[Some("2018-11-03T03:00:00Z"), Some("2018-11-04T03:00:00Z")], + ); + } + + #[test] + fn test_timestamp_trunc_scalar_format_dictionary() { + let mut builder = PrimitiveDictionaryBuilder::::new(); + builder + .append(instant_micros("2024-05-17T12:34:56Z")) + .unwrap(); + builder + .append(instant_micros("2024-06-30T23:59:59Z")) + .unwrap(); + builder + .append(instant_micros("2024-05-17T12:34:56Z")) + .unwrap(); + builder.append_null(); + let input = builder.finish(); + let input = input.with_values(Arc::new( + input + .values() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + .with_timezone_utc(), + )); + let input_keys = input.keys().clone(); + + let result = timestamp_trunc_dyn(&input, "MONTH".to_string()).unwrap(); + let result = result + .as_any() + .downcast_ref::>() + .unwrap(); + assert_eq!(result.keys(), &input_keys); + let decoded = result.downcast_dict::().unwrap(); + assert_eq!( + decoded.into_iter().collect::>(), + vec![ + Some(instant_micros("2024-05-01T00:00:00Z")), + Some(instant_micros("2024-06-01T00:00:00Z")), + Some(instant_micros("2024-05-01T00:00:00Z")), + None, + ] + ); } #[test] diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql index 672281627c9..0cb65b5b60d 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_date.sql @@ -15,21 +15,58 @@ -- specific language governing permissions and limitations -- under the License. +-- Dictionary-encoded dates reuse the scalar Date32 path for dictionary values. +-- ConfigMatrix: parquet.enable.dictionary=false,true + statement CREATE TABLE test_trunc_date(d date) USING parquet statement -INSERT INTO test_trunc_date VALUES (date('2024-06-15')), (date('2024-01-01')), (date('2024-12-31')), (NULL) +INSERT INTO test_trunc_date VALUES + (date('2024-05-17')), + -- Monday / Tuesday / Sunday for Monday-based WEEK truncation. + (date('2024-05-13')), + (date('2024-05-14')), + (date('2024-05-19')), + (date('2024-01-01')), + (date('2023-12-31')), + (date('2024-03-31')), + (date('2024-04-01')), + (date('2024-06-30')), + (date('2024-07-01')), + (date('2024-10-01')), + (date('2024-02-29')), + (date('2000-02-29')), + (date('1900-02-28')), + (date('1969-12-31')), + (date('1960-02-29')), + (date('1900-01-01')), + -- Valid Spark Date32 outside TimestampNanosecond's range. + (date('3333-05-17')), + (NULL) + +query +SELECT d, trunc(d, 'YEAR'), trunc(d, 'YYYY'), trunc(d, 'YY') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'QUARTER') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'MONTH'), trunc(d, 'MON'), trunc(d, 'MM') FROM test_trunc_date ORDER BY d + +query +SELECT d, trunc(d, 'WEEK') FROM test_trunc_date ORDER BY d query -SELECT trunc(d, 'year') FROM test_trunc_date +SELECT d, trunc(d, 'year'), trunc(d, 'Year'), trunc(d, 'yEaR'), trunc(d, 'month'), trunc(d, 'Mon'), trunc(d, 'week') FROM test_trunc_date ORDER BY d +-- Unsupported scalar formats fall back to Spark's codegen dispatcher and remain NULL-compatible. query -SELECT trunc(d, 'month') FROM test_trunc_date +SELECT d, trunc(d, 'DAY'), trunc(d, 'HOUR'), trunc(d, 'SECOND'), trunc(d, 'invalid'), trunc(d, ' YEAR '), trunc(d, '') FROM test_trunc_date ORDER BY d +-- NULL format is Incompatible on the native path, so this uses the codegen dispatcher. query -SELECT trunc(d, 'quarter') FROM test_trunc_date +SELECT d, trunc(d, NULL) FROM test_trunc_date ORDER BY d --- literal arguments query -SELECT trunc(date('2024-06-15'), 'year'), trunc(date('2024-06-15'), 'month'), trunc(date('2024-06-15'), 'quarter') +SELECT trunc(NULL, 'YEAR'), trunc(NULL, NULL) diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql index 1a8bbef9880..fd50c204ab8 100644 --- a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql +++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp.sql @@ -15,26 +15,64 @@ -- specific language governing permissions and limitations -- under the License. --- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true +-- Keep the wide-range fallback fixture independent of far-future JVM/native timezone rules. +-- Config: spark.sql.session.timeZone=UTC +-- Dictionary-encoded timestamps reuse the scalar timestamp path for dictionary values. +-- ConfigMatrix: parquet.enable.dictionary=false,true statement CREATE TABLE test_trunc_ts(ts timestamp) USING parquet statement -INSERT INTO test_trunc_ts VALUES (timestamp('2024-06-15 10:30:45')), (timestamp('2024-01-01 00:00:00')), (NULL) +INSERT INTO test_trunc_ts VALUES + (timestamp('2024-05-17 12:34:56.123456')), + (timestamp('2024-02-29 23:59:59.999999')), + (timestamp('2000-02-29 00:00:00')), + (timestamp('1900-02-28 00:00:00')), + (timestamp('1969-12-31 23:59:59.123456')), + -- Valid Spark timestamp outside TimestampNanosecond's range. + (timestamp('3333-05-17 12:34:56.123456')), + (NULL) query -SELECT date_trunc('year', ts) FROM test_trunc_ts +SELECT ts, date_trunc('YEAR', ts), date_trunc('YYYY', ts), date_trunc('YY', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('month', ts) FROM test_trunc_ts +SELECT ts, date_trunc('QUARTER', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('day', ts) FROM test_trunc_ts +SELECT ts, date_trunc('MONTH', ts), date_trunc('MON', ts), date_trunc('MM', ts) FROM test_trunc_ts ORDER BY ts query -SELECT date_trunc('hour', ts) FROM test_trunc_ts +SELECT ts, date_trunc('WEEK', ts), date_trunc('DAY', ts), date_trunc('DD', ts) FROM test_trunc_ts ORDER BY ts --- literal arguments query -SELECT date_trunc('year', timestamp('2024-06-15 10:30:45')), date_trunc('month', timestamp('2024-06-15 10:30:45')), date_trunc('day', timestamp('2024-06-15 10:30:45')) +SELECT + ts, + date_trunc('HOUR', ts), + date_trunc('MINUTE', ts), + date_trunc('SECOND', ts), + date_trunc('MILLISECOND', ts), + date_trunc('MICROSECOND', ts) +FROM test_trunc_ts +ORDER BY ts + +query +SELECT + ts, + date_trunc('year', ts), + date_trunc('Year', ts), + date_trunc('yEaR', ts), + date_trunc('month', ts), + date_trunc('Mon', ts), + date_trunc('week', ts) +FROM test_trunc_ts +ORDER BY ts + +-- NULL format is Incompatible on the native path. Without allowIncompatible the +-- codegen dispatcher runs Spark's TruncTimestamp and returns NULL. +query +SELECT ts, date_trunc(NULL, ts) FROM test_trunc_ts ORDER BY ts + +query +SELECT date_trunc('YEAR', NULL), date_trunc(NULL, NULL) 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..5817d7a9faa 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 @@ -15,36 +15,49 @@ -- specific language governing permissions and limitations -- under the License. --- 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. +-- Differential coverage for scalar date_trunc around DST overlaps and gaps. +-- Explicit UTC offsets include both occurrences of the repeated US fall-back hour. The 2018 +-- Sao Paulo values exercise its historic midnight spring-forward gap: truncating a valid 01:30 +-- local timestamp to DAY targets the nonexistent local midnight. -- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true --- Config: spark.sql.session.timeZone=America/Los_Angeles +-- ConfigMatrix: spark.sql.session.timeZone=America/Los_Angeles,America/New_York,America/Sao_Paulo statement CREATE TABLE test_trunc_ambiguous(ts timestamp) USING parquet statement INSERT INTO test_trunc_ambiguous VALUES - (timestamp('2024-11-03 01:30:00')) - -query ignore(native panic: chrono returns None for ambiguous local time during DST fall-back) -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) -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) -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) -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) -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) -SELECT ts, date_trunc('YEAR', ts) FROM test_trunc_ambiguous ORDER BY ts + (timestamp('2018-11-04T01:30:15.123456Z')), + (timestamp('2018-11-04T02:30:15.123456Z')), + (timestamp('2018-11-04T03:30:15.123456Z')), + (timestamp('2018-11-04T04:30:15.123456Z')), + (timestamp('2024-03-10T06:30:15.123456Z')), + (timestamp('2024-03-10T07:30:15.123456Z')), + (timestamp('2024-03-10T08:30:15.123456Z')), + (timestamp('2024-03-10T09:30:15.123456Z')), + (timestamp('2024-03-10T10:30:15.123456Z')), + (timestamp('2024-03-10T11:30:15.123456Z')), + (timestamp('2024-11-03T05:30:15.123456Z')), + (timestamp('2024-11-03T06:30:15.123456Z')), + (timestamp('2024-11-03T07:30:15.123456Z')), + (timestamp('2024-11-03T08:30:15.123456Z')), + (timestamp('2024-11-03T09:30:15.123456Z')), + (timestamp('2024-11-03T10:30:15.123456Z')), + (NULL) + +query +SELECT + ts, + date_trunc('YEAR', ts), + date_trunc('QUARTER', ts), + date_trunc('MONTH', ts), + date_trunc('WEEK', ts), + date_trunc('DAY', ts), + date_trunc('HOUR', ts), + date_trunc('MINUTE', ts), + date_trunc('SECOND', ts), + date_trunc('MILLISECOND', ts), + date_trunc('MICROSECOND', ts) +FROM test_trunc_ambiguous +ORDER BY ts diff --git a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala index 9b2b73f97e3..c043b05b86c 100644 --- a/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala @@ -41,6 +41,14 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH private val crossTimezones = Seq("UTC", "America/Los_Angeles", "Europe/London", "Asia/Tokyo") + /** + * Dates around 2024 stay inside DataFusion `date_trunc`'s TimestampNanosecond range + * (approximately 1678 through 2262). The fuzz generator's default `baseDate` is year 3333, + * which only exercises Comet's wide-range fallback. + */ + private val dataFusionRangeBaseDate: Long = + new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime + private def deepestSparkThrowable(error: Throwable): SparkThrowable with Throwable = causeChain(error) .collect { case e: SparkThrowable with Throwable => e } @@ -80,6 +88,14 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } test("trunc (TruncDate)") { + checkTruncDate(dataFusionRangeBaseDate) + } + + test("trunc (TruncDate) - dates outside DataFusion nanosecond range") { + checkTruncDate(FuzzDataGenerator.defaultBaseDate) + } + + private def checkTruncDate(baseDate: Long): Unit = { val supportedFormats = CometTruncDate.supportedFormats val unsupportedFormats = Seq("invalid") @@ -88,7 +104,12 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH Seq( StructField("c0", DataTypes.DateType, true), StructField("c1", DataTypes.StringType, true))) - val df = FuzzDataGenerator.generateDataFrame(r, spark, schema, 1000, DataGenOptions()) + val df = FuzzDataGenerator.generateDataFrame( + r, + spark, + schema, + 1000, + DataGenOptions(baseDate = baseDate)) df.createOrReplaceTempView("tbl") @@ -112,7 +133,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH val supportedFormats = CometTruncTimestamp.supportedFormats val unsupportedFormats = Seq("invalid") - createTimestampTestData().createOrReplaceTempView("tbl") + createTimestampTestData(dataFusionRangeBaseDate).createOrReplaceTempView("tbl") // TODO test fails with non-UTC timezone // https://github.com/apache/datafusion-comet/issues/2649 @@ -138,7 +159,9 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH val unsupportedFormats = Seq("invalid") withTempDir { path => - createTimestampTestData().write.mode(SaveMode.Overwrite).parquet(path.toString) + createTimestampTestData(dataFusionRangeBaseDate).write + .mode(SaveMode.Overwrite) + .parquet(path.toString) spark.read.parquet(path.toString).createOrReplaceTempView("tbl") // TODO test fails with non-UTC timezone @@ -163,6 +186,15 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH } } + test("date_trunc (TruncTimestamp) - timestamps outside DataFusion nanosecond range") { + createTimestampTestData(FuzzDataGenerator.defaultBaseDate).createOrReplaceTempView("tbl") + withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") { + for (format <- Seq("year", "week", "hour")) { + checkSparkAnswerAndOperator(s"SELECT c0, date_trunc('$format', c0) from tbl order by c0") + } + } + } + test("date_trunc - non-UTC timezone routes via codegen dispatcher") { // The native date_trunc is Incompatible in non-UTC sessions // (https://github.com/apache/datafusion-comet/issues/2649), so with allowIncompatible=false @@ -207,9 +239,7 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH // for any non-UTC timezone whose DST transitions are 1-hour and whose dates fall within // chrono-tz's precomputed DST horizon (currently ~year 2100). A base date in 2024 keeps the // generator well inside that window. - val baseDate2024 = - new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime - createTimestampTestData(baseDate2024).createOrReplaceTempView("tbl") + createTimestampTestData(dataFusionRangeBaseDate).createOrReplaceTempView("tbl") val nonUtcTimezones = Seq("America/New_York", "Europe/London", "Asia/Tokyo") for (tz <- nonUtcTimezones) { @@ -293,14 +323,12 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH // Use a reasonable date range (around year 2024) to avoid chrono-tz DST calculation // issues with far-future dates. The default baseDate is year 3333 which is beyond // the range where chrono-tz can reliably calculate DST transitions. - val reasonableBaseDate = - new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime val ntzDF = FuzzDataGenerator.generateDataFrame( r, spark, ntzSchema, 100, - DataGenOptions(baseDate = reasonableBaseDate)) + DataGenOptions(baseDate = dataFusionRangeBaseDate)) ntzDF.createOrReplaceTempView("ntz_tbl") for (tz <- crossTimezones) { withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz) { From 0e85f5a9d717e6c44f9ab3a72173d4ff3c3fa3ce Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Wed, 16 Sep 2026 12:21:50 +0800 Subject: [PATCH 2/3] perf: retain direct Date32 truncation kernel --- native/spark-expr/benches/date_trunc.rs | 36 ++++++-- native/spark-expr/src/kernels/temporal.rs | 101 +--------------------- 2 files changed, 35 insertions(+), 102 deletions(-) diff --git a/native/spark-expr/benches/date_trunc.rs b/native/spark-expr/benches/date_trunc.rs index a3dd0900c61..b8342c247df 100644 --- a/native/spark-expr/benches/date_trunc.rs +++ b/native/spark-expr/benches/date_trunc.rs @@ -18,31 +18,57 @@ use arrow::array::{ArrayRef, Date32Array}; use criterion::{criterion_group, criterion_main, Criterion}; use datafusion_comet_spark_expr::date_trunc_dyn; -use std::sync::Arc; +use std::{hint::black_box, sync::Arc}; fn criterion_benchmark(c: &mut Criterion) { - let date_array = create_date_array(); + let modern_dates: ArrayRef = Arc::new(create_modern_date_array()); + let mixed_range_null_dates: ArrayRef = Arc::new(create_mixed_range_null_date_array()); let mut group = c.benchmark_group("date_trunc"); // Benchmark each truncation format for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] { - let array_ref: ArrayRef = Arc::new(date_array.clone()); group.bench_function(format!("date_trunc_{}", format.to_lowercase()), |b| { - b.iter(|| date_trunc_dyn(&array_ref, format.to_string()).unwrap()); + b.iter(|| { + black_box(date_trunc_dyn(black_box(&modern_dates), format.to_string()).unwrap()) + }); }); + group.bench_function( + format!("date_trunc_{}_mixed_range_null", format.to_lowercase()), + |b| { + b.iter(|| { + black_box( + date_trunc_dyn(black_box(&mixed_range_null_dates), format.to_string()) + .unwrap(), + ) + }); + }, + ); } group.finish(); } -fn create_date_array() -> Date32Array { +fn create_modern_date_array() -> Date32Array { // Create 10000 dates spanning several years (more realistic workload) // Days since Unix epoch: range from 0 (1970-01-01) to ~19000 (2022) let dates: Vec = (0..10000).map(|i| (i * 2) % 19000).collect(); Date32Array::from(dates) } +fn create_mixed_range_null_date_array() -> Date32Array { + // Mix ordinary modern dates with dates around year 3333, which are outside + // DataFusion's TimestampNanosecond range, and nulls. + let dates: Vec> = (0..10000) + .map(|i| match i % 10 { + 0 | 1 => None, + 2 | 3 => Some(497_826 + (i * 17) % 365), + _ => Some((i * 2) % 19000), + }) + .collect(); + Date32Array::from(dates) +} + fn config() -> Criterion { Criterion::default() } diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index b0f5fb200aa..4fe5ee7df66 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -32,7 +32,6 @@ use arrow::array::{ }; use arrow::{ array::*, - compute::cast, datatypes::{DataType, Field, TimeUnit}, }; use datafusion::{ @@ -368,18 +367,6 @@ where /// Truncates a date expressed as days since the epoch, returning `None` if it is out of range. type DateTruncFn = fn(i32) -> Option; -/// The Spark `trunc` spellings and the canonical granularities accepted by DataFusion. -const DATE_TRUNC_ALIASES: [(&str, &str); 8] = [ - ("YEAR", "year"), - ("YYYY", "year"), - ("YY", "year"), - ("QUARTER", "quarter"), - ("MONTH", "month"), - ("MON", "month"), - ("MM", "month"), - ("WEEK", "week"), -]; - /// The `date_trunc` formats Spark accepts, and the truncation each one selects. const DATE_TRUNC_FORMATS: [(&str, DateTruncFn); 8] = [ ("YEAR", trunc_days_to_year), @@ -409,19 +396,6 @@ fn date_trunc_fn_for_format(format: &str) -> Result { }) } -/// Normalize a Spark `trunc` format without exposing additional DataFusion granularities. -fn normalize_date_trunc_format(format: &str) -> Result<&'static str, SparkError> { - DATE_TRUNC_ALIASES - .iter() - .find(|(name, _)| name.eq_ignore_ascii_case(format)) - .map(|(_, granularity)| *granularity) - .ok_or_else(|| { - SparkError::Internal(format!( - "Unsupported format: {format:?} for function 'date_trunc'" - )) - }) -} - const MICROS_PER_DAY: i64 = 86_400_000_000; #[inline] @@ -447,72 +421,13 @@ fn fits_datafusion_coarse_trunc_range(micros: i64) -> bool { && micros >= LOWER_NANOSECOND_MICROS + COARSE_TRUNC_MARGIN_MICROS } -#[inline] -fn date32_to_utc_midnight_micros(days: i32) -> Option { - i64::from(days).checked_mul(MICROS_PER_DAY) -} - -#[inline] -fn date32_fits_upstream(days: i32) -> bool { - date32_to_utc_midnight_micros(days).is_some_and(fits_datafusion_coarse_trunc_range) -} - -/// Truncate scalar-format Date32 values through DataFusion's physical `date_trunc`. +/// Truncate Date32 values directly in days since the epoch. /// -/// DataFusion 55.1 scales coarse timestamp granularities to nanoseconds internally. Spark Date -/// supports approximately years 0001 through 9999, while TimestampNanosecond only spans roughly -/// 1677 through 2262. Values outside the guarded TimestampNanosecond range therefore retain the -/// established Date32 calculation; values inside it use the upstream cast sandwich. +/// Routing Date32 through DataFusion's timestamp kernel requires two casts and temporary arrays, +/// which is materially slower than this single-pass implementation. fn date_trunc_date32(array: &Date32Array, format: String) -> Result { - let granularity = normalize_date_trunc_format(&format)?; let trunc_fn = date_trunc_fn_for_format(&format)?; - let mut has_wide_value = false; - let upstream_input: Date32Array = array - .iter() - .map(|value| { - value.and_then(|days| { - if date32_fits_upstream(days) { - Some(days) - } else { - has_wide_value = true; - None - } - }) - }) - .collect(); - - if upstream_input.null_count() == array.len() { - return Ok(array.iter().map(|value| value.and_then(trunc_fn)).collect()); - } - - let timestamps = cast( - &upstream_input, - &DataType::Timestamp(TimeUnit::Microsecond, None), - )?; - let truncated = datafusion_date_trunc(timestamps, granularity)?; - let truncated = cast(truncated.as_ref(), &DataType::Date32)?; - let upstream = truncated - .as_any() - .downcast_ref::() - .expect("DataFusion date_trunc Date32 cast mismatch"); - - if !has_wide_value { - return Ok(upstream.clone()); - } - - Ok(array - .iter() - .enumerate() - .map(|(index, value)| { - value.and_then(|days| { - if date32_fits_upstream(days) { - Some(upstream.value(index)) - } else { - trunc_fn(days) - } - }) - }) - .collect()) + Ok(array.iter().map(|value| value.and_then(trunc_fn)).collect()) } /// @@ -1367,14 +1282,6 @@ mod tests { format!("Unsupported format: {format:?} for function 'date_trunc'") ); } - - assert!(super::date32_to_utc_midnight_micros(i32::MAX).is_none()); - assert!(super::date32_to_utc_midnight_micros(i32::MIN).is_none()); - assert!(!super::date32_fits_upstream(epoch_days("3333-05-17"))); - assert!(!super::date32_fits_upstream(epoch_days("0001-01-01"))); - assert!(super::date32_fits_upstream(epoch_days("2024-05-17"))); - assert!(super::date32_fits_upstream(epoch_days("1678-09-27"))); - assert!(!super::date32_fits_upstream(epoch_days("1678-09-26"))); } #[test] From 173acf3c357f2d23dd2c34db28c554771ded77bc Mon Sep 17 00:00:00 2001 From: sam-1112 Date: Tue, 29 Sep 2026 16:00:00 +0800 Subject: [PATCH 3/3] fix: preserve timezone semantics for minute truncation --- native/spark-expr/src/kernels/temporal.rs | 111 +++++++++++++++++- .../trunc_timestamp_dst_ambiguous.sql | 15 ++- 2 files changed, 120 insertions(+), 6 deletions(-) diff --git a/native/spark-expr/src/kernels/temporal.rs b/native/spark-expr/src/kernels/temporal.rs index 4fe5ee7df66..f11915a8d4f 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/native/spark-expr/src/kernels/temporal.rs @@ -18,7 +18,8 @@ //! temporal kernels use chrono::{ - DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, TimeZone, Timelike, Utc, + DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, Offset, TimeZone, + Timelike, Utc, }; use std::sync::Arc; @@ -90,6 +91,7 @@ macro_rules! return_compute_error_with { // The number of days between the beginning of the proleptic gregorian calendar (0001-01-01) // and the beginning of the Unix Epoch (1970-01-01) const DAYS_TO_UNIX_EPOCH: i32 = 719_163; +const MICROS_PER_MINUTE: i64 = 60_000_000; // Optimized date truncation functions that work directly with days since epoch // These avoid the overhead of converting to/from NaiveDateTime @@ -811,11 +813,87 @@ fn datafusion_timestamp_trunc_requires_nanos(granularity: &str, has_timezone: bo } } +/// Returns whether a timezone has a zero UTC offset at every instant. +/// +/// Keep this list conservative: an unlisted timezone takes the zone-aware path, which is always +/// correct even when its current offset happens to be zero. +fn is_utc_timezone(timezone: &str) -> bool { + matches!( + timezone, + "UTC" | "Etc/UTC" | "Etc/GMT" | "GMT" | "Z" | "+00:00" | "-00:00" | "00:00" + ) +} + +/// Truncate timezone-aware timestamps to the local minute boundary. +/// +/// DataFusion floors the stored UTC microseconds directly. That is only equivalent to Spark's +/// local-time truncation when the timezone offset is a whole number of minutes. Historical offsets +/// can include seconds, so resolve the offset for each instant before finding the local remainder. +/// If truncation crosses an offset transition, re-resolve the truncated local datetime just like +/// `ZonedDateTime.truncatedTo` rather than retaining the input instant's offset. +fn timestamp_trunc_minute_tz( + array: &TimestampMicrosecondArray, + timezone: &str, +) -> Result { + as_timestamp_tz_with_op::<&TimestampMicrosecondArray, TimestampMicrosecondType, _>( + ArrayIter::new(array), + TimestampMicrosecondBuilder::with_capacity(array.len()), + timezone, + |dt| { + let micros = dt.timestamp_micros(); + let timezone = dt.timezone(); + let original_offset_secs = dt.offset().fix().local_minus_utc(); + let offset_micros = i64::from(original_offset_secs) * 1_000_000; + let candidate = micros - (micros + offset_micros).rem_euclid(MICROS_PER_MINUTE); + let candidate_dt = + as_datetime_with_timezone::(candidate, timezone) + .expect("truncated minute candidate must be a valid datetime"); + let candidate_offset_secs = candidate_dt.offset().fix().local_minus_utc(); + + if candidate_offset_secs == original_offset_secs { + return candidate; + } + + let truncated_local = dt + .naive_local() + .with_second(0) + .and_then(|local| local.with_nanosecond(0)) + .expect("truncated local minute must be a valid datetime"); + match timezone.from_local_datetime(&truncated_local) { + LocalResult::Single(resolved) => resolved.timestamp_micros(), + LocalResult::Ambiguous(earlier, later) => { + // ZonedDateTime retains the original offset when it is valid in an overlap. + if earlier.offset().fix().local_minus_utc() == original_offset_secs { + earlier.timestamp_micros() + } else if later.offset().fix().local_minus_utc() == original_offset_secs { + later.timestamp_micros() + } else { + earlier.timestamp_micros() + } + } + LocalResult::None => { + // The candidate lies immediately before a forward transition. Java advances + // a nonexistent local time by the gap, which is equivalent to resolving it + // with the candidate's pre-transition offset. + naive_to_micros(truncated_local) - i64::from(candidate_offset_secs) * 1_000_000 + } + } + }, + ) +} + fn timestamp_trunc_upstream( array: &TimestampMicrosecondArray, format: &str, ) -> Result { let granularity = normalize_timestamp_trunc_format(format)?; + + if granularity == "minute" { + if let Some(timezone) = array.timezone().filter(|tz| !is_utc_timezone(tz)) { + return timestamp_trunc_minute_tz(array, timezone); + } + } + let requires_nanos = datafusion_timestamp_trunc_requires_nanos(granularity, array.timezone().is_some()); @@ -1565,6 +1643,37 @@ mod tests { ); } + #[test] + fn test_timestamp_trunc_minute_with_historical_offset_seconds() { + // Africa/Monrovia used UTC-00:44:30 in 1960. The input instant is local 10:30:45, + // which Spark truncates to local 10:30:00 (11:14:30 UTC), not a UTC minute boundary. + assert_timestamp_trunc( + "MINUTE", + Some("Africa/Monrovia"), + &[Some("1960-06-15T11:15:15Z"), None], + &[Some("1960-06-15T11:14:30Z"), None], + ); + + // Asia/Aden changed from UTC+03:06:52 to UTC+03:00 within the local 23:53 minute. + // Truncating the post-transition 23:53:30 must re-resolve local 23:53:00 with the + // pre-transition offset, matching ZonedDateTime.truncatedTo. + assert_timestamp_trunc( + "MINUTE", + Some("Asia/Aden"), + &[Some("1947-03-13T20:53:30.123Z")], + &[Some("1947-03-13T20:46:08Z")], + ); + + // Monrovia's 1972 transition skipped local 00:00:00 through 00:44:29. Truncating + // 00:44:45 targets the gap, which ZonedDateTime shifts forward by 44 minutes 30 seconds. + assert_timestamp_trunc( + "MINUTE", + Some("Africa/Monrovia"), + &[Some("1972-01-07T00:44:45Z")], + &[Some("1972-01-07T01:28:30Z")], + ); + } + #[test] fn test_timestamp_trunc_scalar_format_dictionary() { let mut builder = PrimitiveDictionaryBuilder::::new(); 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 5817d7a9faa..eac73c99248 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 @@ -15,19 +15,24 @@ -- specific language governing permissions and limitations -- under the License. --- Differential coverage for scalar date_trunc around DST overlaps and gaps. --- Explicit UTC offsets include both occurrences of the repeated US fall-back hour. The 2018 --- Sao Paulo values exercise its historic midnight spring-forward gap: truncating a valid 01:30 --- local timestamp to DAY targets the nonexistent local midnight. +-- Differential coverage for scalar date_trunc around DST overlaps, gaps, and historical offsets +-- containing seconds. Explicit UTC offsets include both occurrences of the repeated US fall-back +-- hour. The 2018 Sao Paulo values exercise its historic midnight spring-forward gap: truncating a +-- valid 01:30 local timestamp to DAY targets the nonexistent local midnight. Africa/Monrovia used +-- UTC-00:44:30 until 1972, so its local minute boundaries do not align with UTC minute boundaries. +-- Asia/Aden covers minute truncation across a historical offset transition. -- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true --- ConfigMatrix: spark.sql.session.timeZone=America/Los_Angeles,America/New_York,America/Sao_Paulo +-- ConfigMatrix: spark.sql.session.timeZone=America/Los_Angeles,America/New_York,America/Sao_Paulo,Africa/Monrovia,Asia/Aden statement CREATE TABLE test_trunc_ambiguous(ts timestamp) USING parquet statement INSERT INTO test_trunc_ambiguous VALUES + (TIMESTAMP '1960-06-15 10:30:45'), + (timestamp('1947-03-13T20:53:30.123Z')), + (timestamp('1972-01-07T00:44:45Z')), (timestamp('2018-11-04T01:30:15.123456Z')), (timestamp('2018-11-04T02:30:15.123456Z')), (timestamp('2018-11-04T03:30:15.123456Z')),