From 408c5be194473274f2a89375d514c372158895b4 Mon Sep 17 00:00:00 2001 From: naman Date: Sat, 19 Sep 2026 15:56:57 +0530 Subject: [PATCH] fix: truncate scalar timestamps in their own unit in date_trunc A scalar timestamp was always converted to nanoseconds before truncating, so a Timestamp(Second), Timestamp(Millisecond) or Timestamp(Microsecond) value outside the nanosecond range failed with "out of range", while the same value in a column was truncated in its own unit and succeeded. The scalar path now takes the same route as the array path for the granularities that are plain arithmetic in the input's unit. The unit table and the checked truncation are shared between both paths. --- .../functions/src/datetime/date_trunc.rs | 207 ++++++++++++++---- .../test_files/datetime/timestamps.slt | 25 ++- 2 files changed, 188 insertions(+), 44 deletions(-) diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index 2c677213a4733..28bd34a65980c 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -271,14 +271,7 @@ impl ScalarUDFImpl for DateTruncFunc { let parsed_tz = parse_tz(tz_opt)?; let array = as_primitive_array::(array)?; - // fast path for fine granularity - // For modern timezones, it's correct to truncate "minute" in this way. - // Both datafusion and arrow are ignoring historical timezone's non-minute granularity - // bias (e.g., Asia/Kathmandu before 1919 is UTC+05:41:16). - // In UTC, "hour" and "day" have uniform durations and can be truncated with simple arithmetic - if granularity.is_fine_granularity() - || (parsed_tz.is_none() && granularity.is_fine_granularity_utc()) - { + if truncates_in_input_unit(granularity, parsed_tz.as_ref()) { let result = general_date_trunc_array_fine_granularity( T::UNIT, array, @@ -300,10 +293,16 @@ impl ScalarUDFImpl for DateTruncFunc { tz_opt: Option<&Arc>, ) -> Result { let parsed_tz = parse_tz(tz_opt)?; - let value = if let Some(v) = v { - Some(general_date_trunc(T::UNIT, *v, parsed_tz, granularity)?) - } else { - None + let value = match v { + // Truncate in the input's own unit, as `process_array` does, so a + // scalar accepts every timestamp a column accepts. + // `general_date_trunc` converts to nanoseconds first and so rejects + // values outside the nanosecond range. + Some(v) if truncates_in_input_unit(granularity, parsed_tz.as_ref()) => { + Some(date_trunc_fine_granularity(T::UNIT, *v, granularity)?) + } + Some(v) => Some(general_date_trunc(T::UNIT, *v, parsed_tz, granularity)?), + None => None, }; let value = ScalarValue::new_timestamp::(value, tz_opt.cloned()); Ok(ColumnarValue::Scalar(value)) @@ -755,30 +754,7 @@ fn general_date_trunc_array_fine_granularity( granularity: DatePart, tz_opt: Option>, ) -> Result { - let unit = match (tu, granularity) { - (Second, DatePart::Minute) => NonZeroI64::new(60), - (Second, DatePart::Hour) => NonZeroI64::new(3600), - (Second, DatePart::Day) => NonZeroI64::new(86400), - - (Millisecond, DatePart::Second) => NonZeroI64::new(1_000), - (Millisecond, DatePart::Minute) => NonZeroI64::new(60_000), - (Millisecond, DatePart::Hour) => NonZeroI64::new(3_600_000), - (Millisecond, DatePart::Day) => NonZeroI64::new(86_400_000), - - (Microsecond, DatePart::Millisecond) => NonZeroI64::new(1_000), - (Microsecond, DatePart::Second) => NonZeroI64::new(1_000_000), - (Microsecond, DatePart::Minute) => NonZeroI64::new(60_000_000), - (Microsecond, DatePart::Hour) => NonZeroI64::new(3_600_000_000), - (Microsecond, DatePart::Day) => NonZeroI64::new(86_400_000_000), - - (Nanosecond, DatePart::Microsecond) => NonZeroI64::new(1_000), - (Nanosecond, DatePart::Millisecond) => NonZeroI64::new(1_000_000), - (Nanosecond, DatePart::Second) => NonZeroI64::new(1_000_000_000), - (Nanosecond, DatePart::Minute) => NonZeroI64::new(60_000_000_000), - (Nanosecond, DatePart::Hour) => NonZeroI64::new(3_600_000_000_000), - (Nanosecond, DatePart::Day) => NonZeroI64::new(86_400_000_000_000), - _ => None, - }; + let unit = fine_granularity_unit(tu, granularity); if let Some(unit) = unit { let unit = unit.get(); @@ -796,13 +772,8 @@ fn general_date_trunc_array_fine_granularity( }) .collect(); let array: PrimitiveArray = if maybe_underflow { - array.try_unary(|value| { - value.checked_sub(value.rem_euclid(unit)).ok_or_else(|| { - exec_datafusion_err!( - "Timestamp {value} out of range after truncating to {granularity}" - ) - }) - })? + array + .try_unary(|value| date_trunc_fine_granularity(tu, value, granularity))? } else { PrimitiveArray::new(values.into(), array.nulls().cloned()) } @@ -814,6 +785,66 @@ fn general_date_trunc_array_fine_granularity( } } +/// Whether truncating to `granularity` is plain arithmetic in the input's own +/// time unit, which [`fine_granularity_unit`] and [`date_trunc_fine_granularity`] +/// then perform. +/// +/// For modern timezones, it's correct to truncate "minute" in this way. +/// Both datafusion and arrow are ignoring historical timezone's non-minute granularity +/// bias (e.g., Asia/Kathmandu before 1919 is UTC+05:41:16). +/// In UTC, "hour" and "day" have uniform durations and can be truncated with simple arithmetic +fn truncates_in_input_unit(granularity: DatePart, tz: Option<&Tz>) -> bool { + granularity.is_fine_granularity() + || (tz.is_none() && granularity.is_fine_granularity_utc()) +} + +/// The length of `granularity` in `tu`, or `None` when `granularity` is no +/// coarser than `tu`, so a value in `tu` is already truncated to it. +fn fine_granularity_unit(tu: TimeUnit, granularity: DatePart) -> Option { + match (tu, granularity) { + (Second, DatePart::Minute) => NonZeroI64::new(60), + (Second, DatePart::Hour) => NonZeroI64::new(3600), + (Second, DatePart::Day) => NonZeroI64::new(86400), + + (Millisecond, DatePart::Second) => NonZeroI64::new(1_000), + (Millisecond, DatePart::Minute) => NonZeroI64::new(60_000), + (Millisecond, DatePart::Hour) => NonZeroI64::new(3_600_000), + (Millisecond, DatePart::Day) => NonZeroI64::new(86_400_000), + + (Microsecond, DatePart::Millisecond) => NonZeroI64::new(1_000), + (Microsecond, DatePart::Second) => NonZeroI64::new(1_000_000), + (Microsecond, DatePart::Minute) => NonZeroI64::new(60_000_000), + (Microsecond, DatePart::Hour) => NonZeroI64::new(3_600_000_000), + (Microsecond, DatePart::Day) => NonZeroI64::new(86_400_000_000), + + (Nanosecond, DatePart::Microsecond) => NonZeroI64::new(1_000), + (Nanosecond, DatePart::Millisecond) => NonZeroI64::new(1_000_000), + (Nanosecond, DatePart::Second) => NonZeroI64::new(1_000_000_000), + (Nanosecond, DatePart::Minute) => NonZeroI64::new(60_000_000_000), + (Nanosecond, DatePart::Hour) => NonZeroI64::new(3_600_000_000_000), + (Nanosecond, DatePart::Day) => NonZeroI64::new(86_400_000_000_000), + _ => None, + } +} + +/// Truncates a single `value` in `tu` to a granularity for which +/// [`truncates_in_input_unit`] holds, without leaving `tu`. +fn date_trunc_fine_granularity( + tu: TimeUnit, + value: i64, + granularity: DatePart, +) -> Result { + let Some(unit) = fine_granularity_unit(tu, granularity) else { + return Ok(value); + }; + let unit = unit.get(); + value.checked_sub(value.rem_euclid(unit)).ok_or_else(|| { + exec_datafusion_err!( + "Timestamp {value} out of range after truncating to {granularity}" + ) + }) +} + // truncates a single value with the given timeunit to the specified granularity fn general_date_trunc( tu: TimeUnit, @@ -1441,6 +1472,96 @@ mod tests { } } + /// Evaluates `date_trunc(granularity, value)` once on a scalar and once on a + /// one-row column holding the same value. + fn date_trunc_scalar_and_array( + granularity: &str, + value: ScalarValue, + ) -> ( + datafusion_common::Result, + datafusion_common::Result, + ) { + let invoke = |arg: ColumnarValue| { + let data_type = value.data_type(); + let args = ScalarFunctionArgs { + args: vec![ColumnarValue::Scalar(ScalarValue::from(granularity)), arg], + arg_fields: vec![ + Field::new("a", DataType::Utf8, false).into(), + Field::new("b", data_type.clone(), true).into(), + ], + number_rows: 1, + return_field: Field::new("f", data_type, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + match DateTruncFunc::new().invoke_with_args(args)? { + ColumnarValue::Scalar(result) => Ok(result), + ColumnarValue::Array(result) => ScalarValue::try_from_array(&result, 0), + } + }; + ( + invoke(ColumnarValue::Scalar(value.clone())), + invoke(ColumnarValue::Array(value.to_array().unwrap())), + ) + } + + /// A timestamp beyond the nanosecond range is truncated the same way as a + /// scalar and as a column: neither converts it to nanoseconds first. + #[test] + fn scalar_and_array_accept_timestamps_beyond_nanosecond_range() { + // 2286-11-20T17:46:40, after the last nanosecond timestamp in 2262 + let seconds = 10_000_000_000; + let utc: Option> = Some("UTC".into()); + let cases = [ + ( + ScalarValue::TimestampSecond(Some(seconds + 59), None), + "minute", + ), + ( + ScalarValue::TimestampSecond(Some(seconds + 1), None), + "hour", + ), + (ScalarValue::TimestampSecond(Some(seconds + 1), None), "day"), + (ScalarValue::TimestampSecond(Some(seconds), None), "second"), + ( + ScalarValue::TimestampSecond(Some(seconds + 59), utc.clone()), + "minute", + ), + ( + ScalarValue::TimestampMillisecond(Some(seconds * 1_000 + 999), None), + "second", + ), + ( + ScalarValue::TimestampMicrosecond(Some(seconds * 1_000_000 + 999), None), + "millisecond", + ), + ( + ScalarValue::TimestampMicrosecond(Some(-seconds * 1_000_000 - 1), utc), + "second", + ), + ]; + for (value, granularity) in cases { + let (scalar, array) = date_trunc_scalar_and_array(granularity, value.clone()); + let scalar = scalar.unwrap_or_else(|e| { + panic!("scalar date_trunc('{granularity}', {value:?}) failed: {e}") + }); + assert_eq!( + scalar, + array.unwrap(), + "date_trunc('{granularity}', {value:?})" + ); + } + + // The issue's example: `to_timestamp_seconds(10000000000)` + let (scalar, _) = date_trunc_scalar_and_array( + "second", + ScalarValue::TimestampSecond(Some(seconds), None), + ); + assert_eq!( + scalar.unwrap(), + ScalarValue::TimestampSecond(Some(seconds), None) + ); + } + fn assert_fine_granularity_underflow( array: PrimitiveArray, granularity: DatePart, diff --git a/datafusion/sqllogictest/test_files/datetime/timestamps.slt b/datafusion/sqllogictest/test_files/datetime/timestamps.slt index 05622c22dad87..2680173ac95ea 100644 --- a/datafusion/sqllogictest/test_files/datetime/timestamps.slt +++ b/datafusion/sqllogictest/test_files/datetime/timestamps.slt @@ -2523,9 +2523,32 @@ SELECT arrow_typeof(date_trunc('hour', TIME '14:30:45')); ---- Time64(ns) +# date_trunc accepts the same timestamps as a scalar as it does as a column, +# including ones beyond the range of a nanosecond timestamp +query P +SELECT date_trunc('second', to_timestamp_seconds(10000000000)); +---- +2286-11-20T17:46:40 + +query P +SELECT date_trunc('second', to_timestamp_seconds(ts)) +FROM (VALUES (10000000000)) AS timestamps(ts); +---- +2286-11-20T17:46:40 + +query P +SELECT date_trunc('minute', to_timestamp_millis(10000000019999)); +---- +2286-11-20T17:46:00 + +query P +SELECT date_trunc('millisecond', to_timestamp_micros(10000000000123456)); +---- +2286-11-20T17:46:40.123 + query error DataFusion error: Execution error: Timestamp 9223372036854775807 out of range SELECT date_trunc( - 'hour', + 'week', arrow_cast(9223372036854775807, 'Timestamp(Second, None)') );