perf: optimize next_day in datafusion-comet-spark-expr

native/spark-expr/Cargo.toml

@@ -135,3 +135,7 @@ harness = false
 [[bench]]
 name = "parse_url"
 harness = false
+
+[[bench]]
+name = "next_day"
+harness = false

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<i32> = (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<ColumnarValue>, 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);

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<Weekday> {
+    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<Weekday> {
-    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]