From b1957b57365ddb0838bd57546d6376a98965a371 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 12 Jul 2026 09:09:09 -0600 Subject: [PATCH] perf: optimize next_day in datafusion-comet-spark-expr --- native/spark-expr/Cargo.toml | 4 + native/spark-expr/benches/next_day.rs | 98 +++++++++++++++++++ .../spark-expr/src/datetime_funcs/next_day.rs | 52 ++++++++-- 3 files changed, 145 insertions(+), 9 deletions(-) create mode 100644 native/spark-expr/benches/next_day.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index ef00bef1ea..50515d9370 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -135,3 +135,7 @@ harness = false [[bench]] name = "parse_url" harness = false + +[[bench]] +name = "next_day" +harness = false diff --git a/native/spark-expr/benches/next_day.rs b/native/spark-expr/benches/next_day.rs new file mode 100644 index 0000000000..9439a5ce6b --- /dev/null +++ b/native/spark-expr/benches/next_day.rs @@ -0,0 +1,98 @@ +// 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. + +//! Benchmarks for the Spark-compatible `next_day` expression. + +use arrow::array::{ArrayRef, Date32Array, StringArray}; +use arrow::datatypes::{DataType, Field}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::config::ConfigOptions; +use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_comet_spark_expr::SparkNextDay; +use std::hint::black_box; +use std::sync::Arc; + +const BATCH_SIZE: usize = 8192; + +fn date_array() -> ArrayRef { + // A spread of valid Date32 values (days since epoch) around the year 2024. + let values: Vec = (0..BATCH_SIZE as i32).map(|i| 19_000 + i % 400).collect(); + Arc::new(Date32Array::from(values)) +} + +/// Column of day-of-week names cycling through recognized names in varied casing. +fn day_of_week_array() -> ArrayRef { + let names = [ + "Mon", + "TUESDAY", + "wed", + "Th", + "friday", + "SA", + "sun", + "Monday", + "TUE", + "Wednesday", + ]; + let values: Vec<&str> = (0..BATCH_SIZE).map(|i| names[i % names.len()]).collect(); + Arc::new(StringArray::from(values)) +} + +fn invoke(udf: &SparkNextDay, args: Vec, number_rows: usize) -> ColumnarValue { + let return_field = Arc::new(Field::new("next_day", DataType::Date32, true)); + udf.invoke_with_args(ScalarFunctionArgs { + args, + number_rows, + return_field, + arg_fields: vec![], + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() +} + +fn criterion_benchmark(c: &mut Criterion) { + let udf = SparkNextDay::new(false); + let dates = date_array(); + let dows = day_of_week_array(); + + // Scalar day-of-week literal (the most common usage, e.g. next_day(d, 'Sunday')). + c.bench_function("next_day: scalar day-of-week", |b| { + b.iter(|| { + let args = vec![ + ColumnarValue::Array(Arc::clone(&dates)), + ColumnarValue::Scalar(datafusion::common::ScalarValue::Utf8(Some( + "Sunday".to_string(), + ))), + ]; + black_box(invoke(&udf, args, BATCH_SIZE)) + }) + }); + + // Array day-of-week column with mixed casing. + c.bench_function("next_day: array day-of-week", |b| { + b.iter(|| { + let args = vec![ + ColumnarValue::Array(Arc::clone(&dates)), + ColumnarValue::Array(Arc::clone(&dows)), + ]; + black_box(invoke(&udf, args, BATCH_SIZE)) + }) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/datetime_funcs/next_day.rs b/native/spark-expr/src/datetime_funcs/next_day.rs index df4c2f9096..310056fdef 100644 --- a/native/spark-expr/src/datetime_funcs/next_day.rs +++ b/native/spark-expr/src/datetime_funcs/next_day.rs @@ -54,19 +54,45 @@ impl Default for SparkNextDay { } } +/// Longest recognized (upper-cased) day-of-week name is "WEDNESDAY" at 9 bytes. +const MAX_DAY_OF_WEEK_LEN: usize = 9; + +/// Match an already upper-cased day-of-week name (as raw bytes) to a [`Weekday`]. +#[inline] +fn weekday_from_uppercase(upper: &[u8]) -> Option { + match upper { + b"SU" | b"SUN" | b"SUNDAY" => Some(Weekday::Sun), + b"MO" | b"MON" | b"MONDAY" => Some(Weekday::Mon), + b"TU" | b"TUE" | b"TUESDAY" => Some(Weekday::Tue), + b"WE" | b"WED" | b"WEDNESDAY" => Some(Weekday::Wed), + b"TH" | b"THU" | b"THURSDAY" => Some(Weekday::Thu), + b"FR" | b"FRI" | b"FRIDAY" => Some(Weekday::Fri), + b"SA" | b"SAT" | b"SATURDAY" => Some(Weekday::Sat), + _ => None, + } +} + /// Match a day-of-week name to a [`Weekday`]. Mirrors Spark's /// `DateTimeUtils.getDayOfWeekFromString`: case-insensitive, but with no whitespace trimming. fn day_of_week_from_string(day_of_week: &str) -> Option { - match day_of_week.to_uppercase().as_str() { - "SU" | "SUN" | "SUNDAY" => Some(Weekday::Sun), - "MO" | "MON" | "MONDAY" => Some(Weekday::Mon), - "TU" | "TUE" | "TUESDAY" => Some(Weekday::Tue), - "WE" | "WED" | "WEDNESDAY" => Some(Weekday::Wed), - "TH" | "THU" | "THURSDAY" => Some(Weekday::Thu), - "FR" | "FRI" | "FRIDAY" => Some(Weekday::Fri), - "SA" | "SAT" | "SATURDAY" => Some(Weekday::Sat), - _ => None, + // Fast path: for ASCII input (the overwhelmingly common case) Unicode + // upper-casing is identical to ASCII upper-casing, so upper-case into a + // stack buffer and avoid the per-row heap allocation that `to_uppercase` + // incurs. + if day_of_week.is_ascii() { + let bytes = day_of_week.as_bytes(); + // The buffer only needs to hold the longest recognized name; anything + // longer cannot match, so reject it rather than overflow the buffer. + let mut buf = [0u8; MAX_DAY_OF_WEEK_LEN]; + if bytes.len() > buf.len() { + return None; + } + let upper = &mut buf[..bytes.len()]; + upper.copy_from_slice(bytes); + upper.make_ascii_uppercase(); + return weekday_from_uppercase(upper); } + weekday_from_uppercase(day_of_week.to_uppercase().as_bytes()) } /// The first date strictly after `days` (days since the Unix epoch) that falls on `weekday`. @@ -174,6 +200,14 @@ mod tests { assert_eq!(day_of_week_from_string("MO "), None); assert_eq!(day_of_week_from_string(""), None); assert_eq!(day_of_week_from_string("NOT_A_DAY"), None); + // Mixed case is accepted on the ASCII fast path. + assert_eq!(day_of_week_from_string("wEdNeSdAy"), Some(Weekday::Wed)); + // Inputs longer than any recognized name never match. + assert_eq!(day_of_week_from_string("SUNDAYSUNDAY"), None); + // Non-ASCII input falls back to full Unicode upper-casing, matching the + // reference behavior (e.g. U+017F LONG S upper-cases to 'S'). + assert_eq!(day_of_week_from_string("\u{17f}unday"), Some(Weekday::Sun)); + assert_eq!(day_of_week_from_string("úñîçödé"), None); } #[test]