perf: optimize date_trunc (>2x faster)

native/spark-expr/Cargo.toml

@@ -139,3 +139,7 @@ harness = false
 [[bench]]
 name = "to_json"
 harness = false
+
+[[bench]]
+name = "date_trunc_array_fmt"
+harness = false

native/spark-expr/benches/date_trunc_array_fmt.rs

@@ -0,0 +1,91 @@
+// 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.
+
+use arrow::array::{ArrayRef, Date32Array, StringArray};
+use arrow::datatypes::{DataType, Field};
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::common::config::ConfigOptions;
+use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
+use datafusion_comet_spark_expr::SparkDateTrunc;
+use std::hint::black_box;
+use std::sync::Arc;
+
+const SIZE: usize = 8192;
+
+fn create_date_array() -> ArrayRef {
+    let dates: Vec<i32> = (0..SIZE).map(|i| ((i * 2) % 19000) as i32).collect();
+    Arc::new(Date32Array::from(dates))
+}
+
+/// A constant format column, which is what a query like `trunc(d, fmt)` produces when the
+/// format comes from a column rather than a literal.
+fn create_constant_format_array(format: &str) -> ArrayRef {
+    Arc::new(StringArray::from(vec![format; SIZE]))
+}
+
+/// A low-cardinality format column mixing the supported formats.
+fn create_mixed_format_array() -> ArrayRef {
+    let formats = ["YEAR", "quarter", "MONTH", "week"];
+    let values: Vec<&str> = (0..SIZE).map(|i| formats[i % formats.len()]).collect();
+    Arc::new(StringArray::from(values))
+}
+
+fn invoke(udf: &SparkDateTrunc, dates: &ArrayRef, formats: &ArrayRef) -> ColumnarValue {
+    let args = ScalarFunctionArgs {
+        args: vec![
+            ColumnarValue::Array(Arc::clone(dates)),
+            ColumnarValue::Array(Arc::clone(formats)),
+        ],
+        number_rows: SIZE,
+        return_field: Arc::new(Field::new("date_trunc", DataType::Date32, true)),
+        arg_fields: vec![],
+        config_options: Arc::new(ConfigOptions::default()),
+    };
+    udf.invoke_with_args(args).unwrap()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let udf = SparkDateTrunc::new();
+    let dates = create_date_array();
+
+    let mut group = c.benchmark_group("date_trunc_array_fmt");
+
+    for format in ["YEAR", "QUARTER", "MONTH", "WEEK"] {
+        let formats = create_constant_format_array(format);
+        group.bench_function(format.to_lowercase(), |b| {
+            b.iter(|| black_box(invoke(&udf, &dates, &formats)))
+        });
+    }
+
+    let formats = create_mixed_format_array();
+    group.bench_function("mixed", |b| {
+        b.iter(|| black_box(invoke(&udf, &dates, &formats)))
+    });
+
+    group.finish();
+}
+
+fn config() -> Criterion {
+    Criterion::default()
+}
+
+criterion_group! {
+    name = benches;
+    config = config();
+    targets = criterion_benchmark
+}
+criterion_main!(benches);

native/spark-expr/src/kernels/temporal.rs

@@ -21,6 +21,7 @@ use chrono::{
     DateTime, Datelike, Duration, LocalResult, NaiveDate, NaiveDateTime, TimeZone, Timelike, Utc,
 };
 
+use std::borrow::Cow;
 use std::sync::Arc;
 
 use arrow::array::{
@@ -321,21 +322,48 @@ where
     }
 }
 
+/// Truncates a date expressed as days since the epoch, returning `None` if it is out of range.
+type DateTruncFn = fn(i32) -> Option<i32>;
+
+/// The `date_trunc` formats Spark accepts, and the truncation each one selects.
+const DATE_TRUNC_FORMATS: [(&str, DateTruncFn); 8] = [
+    ("YEAR", trunc_days_to_year),
+    ("YYYY", trunc_days_to_year),
+    ("YY", trunc_days_to_year),
+    ("QUARTER", trunc_days_to_quarter),
+    ("MONTH", trunc_days_to_month),
+    ("MON", trunc_days_to_month),
+    ("MM", trunc_days_to_month),
+    ("WEEK", trunc_days_to_week),
+];
+
+/// Resolve a `date_trunc` format string to the corresponding truncation function.
+///
+/// The format is matched case-insensitively. Every supported format is ASCII, so ASCII input can
+/// be compared in place, leaving the per-row format column path allocation-free. Non-ASCII input
+/// still needs full Unicode case folding to fold the same way Spark does.
+fn date_trunc_fn_for_format(format: &str) -> Result<DateTruncFn, SparkError> {
+    let key: Cow<str> = if format.is_ascii() {
+        Cow::Borrowed(format)
+    } else {
+        Cow::Owned(format.to_uppercase())
+    };
+    DATE_TRUNC_FORMATS
+        .iter()
+        .find(|(name, _)| name.eq_ignore_ascii_case(&key))
+        .map(|(_, trunc_fn)| *trunc_fn)
+        .ok_or_else(|| {
+            SparkError::Internal(format!(
+                "Unsupported format: {format:?} for function 'date_trunc'"
+            ))
+        })
+}
+
 /// Optimized date truncation for Date32 arrays
 /// Works directly with days since epoch instead of converting to/from NaiveDateTime
 fn date_trunc_date32(array: &Date32Array, format: String) -> Result<Date32Array, SparkError> {
     // Select the truncation function based on format
-    let trunc_fn: fn(i32) -> Option<i32> = match format.to_uppercase().as_str() {
-        "YEAR" | "YYYY" | "YY" => trunc_days_to_year,
-        "QUARTER" => trunc_days_to_quarter,
-        "MONTH" | "MON" | "MM" => trunc_days_to_month,
-        "WEEK" => trunc_days_to_week,
-        _ => {
-            return Err(SparkError::Internal(format!(
-                "Unsupported format: {format:?} for function 'date_trunc'"
-            )))
-        }
-    };
+    let trunc_fn = date_trunc_fn_for_format(&format)?;
 
     // Apply truncation to each element
     let result: Date32Array = array
@@ -438,20 +466,19 @@ macro_rules! date_trunc_array_fmt_helper {
         let iter = $array.into_iter();
         match $datatype {
             DataType::Date32 => {
+                // Format columns are almost always constant or very low cardinality, so remember
+                // the last format seen and skip re-resolving it for every row.
+                let mut cached: Option<(&str, DateTruncFn)> = None;
                 for (index, val) in iter.enumerate() {
-                    let trunc_fn: fn(i32) -> Option<i32> =
-                        match $formats.value(index).to_uppercase().as_str() {
-                            "YEAR" | "YYYY" | "YY" => trunc_days_to_year,
-                            "QUARTER" => trunc_days_to_quarter,
-                            "MONTH" | "MON" | "MM" => trunc_days_to_month,
-                            "WEEK" => trunc_days_to_week,
-                            _ => {
-                                return Err(SparkError::Internal(format!(
-                                    "Unsupported format: {:?} for function 'date_trunc'",
-                                    $formats.value(index)
-                                )))
-                            }
-                        };
+                    let format = $formats.value(index);
+                    let trunc_fn = match cached {
+                        Some((cached_format, trunc_fn)) if cached_format == format => trunc_fn,
+                        _ => {
+                            let trunc_fn = date_trunc_fn_for_format(format)?;
+                            cached = Some((format, trunc_fn));
+                            trunc_fn
+                        }
+                    };
                     match val.and_then(trunc_fn) {
                         Some(days) => builder.append_value(days),
                         None => builder.append_null(),