fix: truncate scalar timestamps in their own unit in date_trunc

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::<T>(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,22 @@ impl ScalarUDFImpl for DateTruncFunc {
             tz_opt: Option<&Arc<str>>,
         ) -> Result<ColumnarValue> {
             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()) => {
+                    match fine_granularity_unit(T::UNIT, granularity) {
+                        Some(unit) => {
+                            Some(truncate_to_unit(*v, unit.get(), granularity)?)
+                        }
+                        // `granularity` is no coarser than the input's unit
+                        None => Some(*v),
+                    }
+                }
+                Some(v) => Some(general_date_trunc(T::UNIT, *v, parsed_tz, granularity)?),
+                None => None,
             };
             let value = ScalarValue::new_timestamp::<T>(value, tz_opt.cloned());
             Ok(ColumnarValue::Scalar(value))
@@ -755,30 +760,7 @@ fn general_date_trunc_array_fine_granularity<T: ArrowTimestampType>(
     granularity: DatePart,
     tz_opt: Option<Arc<str>>,
 ) -> Result<ArrayRef> {
-    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 +778,7 @@ fn general_date_trunc_array_fine_granularity<T: ArrowTimestampType>(
             })
             .collect();
         let array: PrimitiveArray<T> = 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| truncate_to_unit(value, unit, granularity))?
         } else {
             PrimitiveArray::new(values.into(), array.nulls().cloned())
         }
@@ -814,6 +790,57 @@ fn general_date_trunc_array_fine_granularity<T: ArrowTimestampType>(
     }
 }
 
+/// Whether truncating to `granularity` is plain arithmetic in the input's own
+/// time unit: rounding down to a multiple of [`fine_granularity_unit`].
+///
+/// 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<NonZeroI64> {
+    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,
+    }
+}
+
+/// Rounds `value` down to a multiple of `unit`, the length of `granularity` in
+/// the value's time unit, or errors if that falls below `i64::MIN`.
+fn truncate_to_unit(value: i64, unit: i64, granularity: DatePart) -> Result<i64> {
+    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 +1468,116 @@ 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<ScalarValue>,
+        datafusion_common::Result<ScalarValue>,
+    ) {
+        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<Arc<str>> = Some("UTC".into());
+        let cases = [
+            // Seconds, truncated to a minute, an hour and a day
+            (
+                ScalarValue::TimestampSecond(Some(seconds + 59), None),
+                "minute",
+            ),
+            (
+                ScalarValue::TimestampSecond(Some(seconds + 1), None),
+                "hour",
+            ),
+            (ScalarValue::TimestampSecond(Some(seconds + 1), None), "day"),
+            // The granularity is the input's own unit
+            (ScalarValue::TimestampSecond(Some(seconds), None), "second"),
+            // With a time zone
+            (
+                ScalarValue::TimestampSecond(Some(seconds + 59), utc.clone()),
+                "minute",
+            ),
+            // Milliseconds and microseconds, truncated to a coarser unit
+            (
+                ScalarValue::TimestampMillisecond(Some(seconds * 1_000 + 999), None),
+                "second",
+            ),
+            (
+                ScalarValue::TimestampMicrosecond(Some(seconds * 1_000_000 + 999), None),
+                "millisecond",
+            ),
+            // Before the epoch, with a time zone
+            (
+                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:?})"
+            );
+        }
+
+        // `date_trunc('second', to_timestamp_seconds(10000000000))` returns its input
+        // rather than an out of range error.
+        // See <https://github.com/apache/datafusion/issues/25432>.
+        let (scalar, _) = date_trunc_scalar_and_array(
+            "second",
+            ScalarValue::TimestampSecond(Some(seconds), None),
+        );
+        assert_eq!(
+            scalar.unwrap(),
+            ScalarValue::TimestampSecond(Some(seconds), None)
+        );
+    }
+
+    /// The scalar fast path shares the same underflow check as the array
+    /// path: truncating a value within one unit of `i64::MIN` to a coarser
+    /// granularity in its own unit must error rather than wrap around.
+    #[test]
+    fn scalar_and_array_reject_fine_granularity_underflow() {
+        let (scalar, array) = date_trunc_scalar_and_array(
+            "minute",
+            ScalarValue::TimestampSecond(Some(i64::MIN), None),
+        );
+        assert!(scalar.is_err(), "expected scalar path to reject underflow");
+        assert!(array.is_err(), "expected array path to reject underflow");
+    }
+
     fn assert_fine_granularity_underflow<T: ArrowTimestampType>(
         array: PrimitiveArray<T>,
         granularity: DatePart,

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)')
 );