speedup date_trunc (~7x faster) in some cases

datafusion/functions/src/datetime/date_trunc.rs

@@ -28,7 +28,7 @@ use arrow::array::types::{
     ArrowTimestampType, TimestampMicrosecondType, TimestampMillisecondType,
     TimestampNanosecondType, TimestampSecondType,
 };
-use arrow::array::{Array, PrimitiveArray};
+use arrow::array::{Array, ArrayRef, Int64Array, PrimitiveArray};
 use arrow::datatypes::DataType::{self, Null, Timestamp, Utf8, Utf8View};
 use arrow::datatypes::TimeUnit::{self, Microsecond, Millisecond, Nanosecond, Second};
 use datafusion_common::cast::as_primitive_array;
@@ -60,6 +60,8 @@ use chrono::{
     - hour / HOUR
     - minute / MINUTE
     - second / SECOND
+    - millisecond / MILLISECOND
+    - microsecond / MICROSECOND
 "#
     ),
     argument(
@@ -185,6 +187,21 @@ impl ScalarUDFImpl for DateTruncFunc {
         ) -> Result<ColumnarValue> {
             let parsed_tz = parse_tz(tz_opt)?;
             let array = as_primitive_array::<T>(array)?;
+
+            // fast path for fine granularities
+            if matches!(
+                granularity.as_str(),
+                "second" | "minute" | "millisecond" | "microsecond"
+            ) || (parsed_tz.is_none() && granularity.as_str() == "hour")
+            {
+                let result = general_date_trunc_array_fine_granularity(
+                    T::UNIT,
+                    array,
+                    granularity.as_str(),
+                )?;
+                return Ok(ColumnarValue::Array(result));
+            }
+
             let array: PrimitiveArray<T> = array
                 .try_unary(|x| {
                     general_date_trunc(T::UNIT, x, parsed_tz, granularity.as_str())
@@ -423,6 +440,51 @@ fn date_trunc_coarse(granularity: &str, value: i64, tz: Option<Tz>) -> Result<i6
     Ok(value.unwrap())
 }
 
+/// Fast path for fine granularities (hour and smaller) that can be handled
+/// with simple arithmetic operations without calendar complexity.
+///
+/// This function is timezone-agnostic and should only be used when:
+/// - No timezone is specified in the input, OR
+/// - The granularity is less than hour as hour can be affected by DST transitions in some cases
+fn general_date_trunc_array_fine_granularity<T: ArrowTimestampType>(
+    tu: TimeUnit,
+    array: &PrimitiveArray<T>,
+    granularity: &str,
+) -> Result<ArrayRef> {
+    let unit = match (tu, granularity) {
+        (Second, "minute") => Some(Int64Array::new_scalar(60)),
+        (Second, "hour") => Some(Int64Array::new_scalar(3600)),
+
+        (Millisecond, "second") => Some(Int64Array::new_scalar(1_000)),
+        (Millisecond, "minute") => Some(Int64Array::new_scalar(60_000)),
+        (Millisecond, "hour") => Some(Int64Array::new_scalar(3_600_000)),
+
+        (Microsecond, "millisecond") => Some(Int64Array::new_scalar(1_000)),
+        (Microsecond, "second") => Some(Int64Array::new_scalar(1_000_000)),
+        (Microsecond, "minute") => Some(Int64Array::new_scalar(60_000_000)),
+        (Microsecond, "hour") => Some(Int64Array::new_scalar(3_600_000_000)),
+
+        (Nanosecond, "microsecond") => Some(Int64Array::new_scalar(1_000)),
+        (Nanosecond, "millisecond") => Some(Int64Array::new_scalar(1_000_000)),
+        (Nanosecond, "second") => Some(Int64Array::new_scalar(1_000_000_000)),
+        (Nanosecond, "minute") => Some(Int64Array::new_scalar(60_000_000_000)),
+        (Nanosecond, "hour") => Some(Int64Array::new_scalar(3_600_000_000_000)),
+        _ => None,
+    };
+
+    if let Some(unit) = unit {
+        let original_type = array.data_type();
+        let array = arrow::compute::cast(array, &DataType::Int64)?;
+        let array = arrow::compute::kernels::numeric::div(&array, &unit)?;
+        let array = arrow::compute::kernels::numeric::mul(&array, &unit)?;
+        let array = arrow::compute::cast(&array, original_type)?;
+        Ok(array)
+    } else {
+        // truncate to the same or smaller unit
+        Ok(Arc::new(array.clone()))
+    }
+}
+
 // truncates a single value with the given timeunit to the specified granularity
 fn general_date_trunc(
     tu: TimeUnit,

docs/source/user-guide/sql/scalar_functions.md

@@ -2150,6 +2150,8 @@ date_trunc(precision, expression)
   - hour / HOUR
   - minute / MINUTE
   - second / SECOND
+  - millisecond / MILLISECOND
+  - microsecond / MICROSECOND
 
 - **expression**: Time expression to operate on. Can be a constant, column, or function.