diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6f90181ee0..cea70f77fb 100644 --- a/native/spark-expr/Cargo.toml +++ b/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 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 = (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..6b68e1e28d 100644 --- a/native/spark-expr/src/kernels/temporal.rs +++ b/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; + +/// 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 { + let key: Cow = 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 { // Select the truncation function based on format - let trunc_fn: fn(i32) -> Option = 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 = - 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(),