This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 56405e031f perf: optimize `date_trunc` (>2x faster) (#4915)
56405e031f is described below
commit 56405e031f5b644d187770308831544ba2b5b6eb
Author: Andy Grove <[email protected]>
AuthorDate: Wed Jul 22 21:22:29 2026 -0600
perf: optimize `date_trunc` (>2x faster) (#4915)
* perf: optimize date_trunc in datafusion-comet-spark-expr
* perf: table-ize timestamp_trunc format lookup for NTZ and TZ paths
Applies the DATE_TRUNC_FORMATS table pattern to both timestamp_trunc paths.
ntz_trunc_fn_for_format and tz_trunc_fn_for_format resolve the format via
eq_ignore_ascii_case against ASCII tables, eliminating the per-invocation and
per-row to_uppercase() String allocations from the scalar-format tz path and
the array-format helper. Also drops the unused Cow / is_ascii split from
date_trunc_fn_for_format since eq_ignore_ascii_case only folds ASCII and every
table entry is ASCII, so a non-A [...]
---
native/spark-expr/Cargo.toml | 4 +
native/spark-expr/benches/date_trunc_array_fmt.rs | 91 +++++++
native/spark-expr/src/kernels/temporal.rs | 279 ++++++++++------------
3 files changed, 220 insertions(+), 154 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 5565933488..326dc474e6 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -163,6 +163,10 @@ harness = false
name = "to_json"
harness = false
+[[bench]]
+name = "date_trunc_array_fmt"
+harness = false
+
[[bench]]
name = "cast_decimal_to_string"
harness = false
diff --git a/native/spark-expr/benches/date_trunc_array_fmt.rs
b/native/spark-expr/benches/date_trunc_array_fmt.rs
new file mode 100644
index 0000000000..2bbf9158c5
--- /dev/null
+++ b/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);
diff --git a/native/spark-expr/src/kernels/temporal.rs
b/native/spark-expr/src/kernels/temporal.rs
index 118de75a4e..a7332d2099 100644
--- a/native/spark-expr/src/kernels/temporal.rs
+++ b/native/spark-expr/src/kernels/temporal.rs
@@ -321,21 +321,43 @@ 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.
+///
+/// Every supported format is ASCII, so `eq_ignore_ascii_case` on the input
as-is is exactly
+/// what Spark does: a non-ASCII input cannot match any ASCII table entry
after ASCII case
+/// folding, and Spark also only recognizes those ASCII literals. This is
allocation-free.
+fn date_trunc_fn_for_format(format: &str) -> Result<DateTruncFn, SparkError> {
+ DATE_TRUNC_FORMATS
+ .iter()
+ .find(|(name, _)| name.eq_ignore_ascii_case(format))
+ .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 +460,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(),
@@ -557,25 +578,79 @@ fn naive_to_micros(dt: NaiveDateTime) -> i64 {
dt.and_utc().timestamp_micros()
}
+/// Truncates a `NaiveDateTime`, returning `None` if the result is out of
range.
+type NtzTruncFn = fn(NaiveDateTime) -> Option<NaiveDateTime>;
+
+/// Truncates a `DateTime<Tz>`, returning `None` if the result is out of range.
+type TzTruncFn = fn(DateTime<Tz>) -> Option<DateTime<Tz>>;
+
+/// 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.
+const TIMESTAMP_TRUNC_FORMATS_NTZ: [(&str, NtzTruncFn); 15] = [
+ ("YEAR", trunc_date_to_year),
+ ("YYYY", trunc_date_to_year),
+ ("YY", trunc_date_to_year),
+ ("QUARTER", trunc_date_to_quarter),
+ ("MONTH", trunc_date_to_month),
+ ("MON", trunc_date_to_month),
+ ("MM", trunc_date_to_month),
+ ("WEEK", trunc_date_to_week),
+ ("DAY", trunc_date_to_day),
+ ("DD", trunc_date_to_day),
+ ("HOUR", trunc_date_to_hour),
+ ("MINUTE", trunc_date_to_minute),
+ ("SECOND", trunc_date_to_second),
+ ("MILLISECOND", trunc_date_to_ms),
+ ("MICROSECOND", trunc_date_to_microsec),
+];
+
+/// Same formats as `TIMESTAMP_TRUNC_FORMATS_NTZ`, monomorphized for the
timezone-aware path.
+const TIMESTAMP_TRUNC_FORMATS_TZ: [(&str, TzTruncFn); 15] = [
+ ("YEAR", trunc_date_to_year),
+ ("YYYY", trunc_date_to_year),
+ ("YY", trunc_date_to_year),
+ ("QUARTER", trunc_date_to_quarter),
+ ("MONTH", trunc_date_to_month),
+ ("MON", trunc_date_to_month),
+ ("MM", trunc_date_to_month),
+ ("WEEK", trunc_date_to_week),
+ ("DAY", trunc_date_to_day),
+ ("DD", trunc_date_to_day),
+ ("HOUR", trunc_date_to_hour),
+ ("MINUTE", trunc_date_to_minute),
+ ("SECOND", trunc_date_to_second),
+ ("MILLISECOND", trunc_date_to_ms),
+ ("MICROSECOND", trunc_date_to_microsec),
+];
+
/// Resolve a truncation format string to the corresponding NaiveDateTime
truncation function.
-fn ntz_trunc_fn_for_format(
- format: &str,
-) -> Result<fn(NaiveDateTime) -> Option<NaiveDateTime>, SparkError> {
- match format.to_uppercase().as_str() {
- "YEAR" | "YYYY" | "YY" => Ok(trunc_date_to_year),
- "QUARTER" => Ok(trunc_date_to_quarter),
- "MONTH" | "MON" | "MM" => Ok(trunc_date_to_month),
- "WEEK" => Ok(trunc_date_to_week),
- "DAY" | "DD" => Ok(trunc_date_to_day),
- "HOUR" => Ok(trunc_date_to_hour),
- "MINUTE" => Ok(trunc_date_to_minute),
- "SECOND" => Ok(trunc_date_to_second),
- "MILLISECOND" => Ok(trunc_date_to_ms),
- "MICROSECOND" => Ok(trunc_date_to_microsec),
- _ => Err(SparkError::Internal(format!(
- "Unsupported format: {format:?} for function 'timestamp_trunc'"
- ))),
- }
+///
+/// All supported formats are ASCII, so `eq_ignore_ascii_case` on the input
as-is is exactly what
+/// Spark does and allocation-free.
+fn ntz_trunc_fn_for_format(format: &str) -> Result<NtzTruncFn, SparkError> {
+ TIMESTAMP_TRUNC_FORMATS_NTZ
+ .iter()
+ .find(|(name, _)| name.eq_ignore_ascii_case(format))
+ .map(|(_, trunc_fn)| *trunc_fn)
+ .ok_or_else(|| {
+ SparkError::Internal(format!(
+ "Unsupported format: {format:?} for function 'timestamp_trunc'"
+ ))
+ })
+}
+
+/// Timezone-aware sibling of `ntz_trunc_fn_for_format`.
+fn tz_trunc_fn_for_format(format: &str) -> Result<TzTruncFn, SparkError> {
+ TIMESTAMP_TRUNC_FORMATS_TZ
+ .iter()
+ .find(|(name, _)| name.eq_ignore_ascii_case(format))
+ .map(|(_, trunc_fn)| *trunc_fn)
+ .ok_or_else(|| {
+ SparkError::Internal(format!(
+ "Unsupported format: {format:?} for function 'timestamp_trunc'"
+ ))
+ })
}
/// Truncate a TimestampNTZ array without any timezone conversion.
@@ -644,61 +719,10 @@ where
timestamp_trunc_ntz(array, format)
}
DataType::Timestamp(TimeUnit::Microsecond, Some(tz)) => {
- match format.to_uppercase().as_str() {
- "YEAR" | "YYYY" | "YY" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_year(dt))
- })
- }
- "QUARTER" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_quarter(dt))
- })
- }
- "MONTH" | "MON" | "MM" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_month(dt))
- })
- }
- "WEEK" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_week(dt))
- })
- }
- "DAY" | "DD" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_day(dt))
- })
- }
- "HOUR" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_hour(dt))
- })
- }
- "MINUTE" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_minute(dt))
- })
- }
- "SECOND" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_second(dt))
- })
- }
- "MILLISECOND" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
- as_micros_from_unix_epoch_utc(trunc_date_to_ms(dt))
- })
- }
- "MICROSECOND" => {
- as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(iter,
builder, tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_microsec(dt))
- })
- }
- _ => Err(SparkError::Internal(format!(
- "Unsupported format: {format:?} for function
'timestamp_trunc'"
- ))),
- }
+ let trunc_fn = tz_trunc_fn_for_format(&format)?;
+ as_timestamp_tz_with_op::<&PrimitiveArray<T>, T, _>(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'",
@@ -802,63 +826,10 @@ macro_rules! timestamp_trunc_array_fmt_helper {
DataType::Timestamp(TimeUnit::Microsecond, Some(tz_str)) => {
let tz: Tz = tz_str.parse()?;
for (index, val) in iter.enumerate() {
- let op_result = match
$formats.value(index).to_uppercase().as_str() {
- "YEAR" | "YYYY" | "YY" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_year(dt))
- })
- }
- "QUARTER" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_quarter(dt))
- })
- }
- "MONTH" | "MON" | "MM" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_month(dt))
- })
- }
- "WEEK" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_week(dt))
- })
- }
- "DAY" | "DD" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_day(dt))
- })
- }
- "HOUR" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_hour(dt))
- })
- }
- "MINUTE" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_minute(dt))
- })
- }
- "SECOND" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_second(dt))
- })
- }
- "MILLISECOND" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_ms(dt))
- })
- }
- "MICROSECOND" => {
- as_timestamp_tz_with_op_single::<T, _>(val, &mut
builder, &tz, |dt| {
-
as_micros_from_unix_epoch_utc(trunc_date_to_microsec(dt))
- })
- }
- _ => Err(SparkError::Internal(format!(
- "Unsupported format: {:?} for function
'timestamp_trunc'",
- $formats.value(index)
- ))),
- };
- op_result?
+ let trunc_fn =
tz_trunc_fn_for_format($formats.value(index))?;
+ as_timestamp_tz_with_op_single::<T, _>(val, &mut builder,
&tz, |dt| {
+ as_micros_from_unix_epoch_utc(trunc_fn(dt))
+ })?;
}
Ok(builder.finish().with_timezone(tz_str.as_ref()))
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
