Skip to content

fix: [spark] honor ANSI mode in next_day, and fix a panic and a missing argument shape #23892

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andygrove wants to merge 2 commits into
base: main
Choose a base branch
from
+514 −154
Open

fix: [spark] honor ANSI mode in next_day, and fix a panic and a missing argument shape #23892

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9b43895
fix: honor ANSI mode in next_day, and fix a missing argument shape
andygrove Jul 25, 2026
9f94893
Merge remote-tracking branch 'apache/main' into pr23892-merge
andygrove Sep 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions datafusion/spark/src/function/datetime/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,11 +126,9 @@ pub mod expr_fn {
"Returns the three-letter abbreviated month name from a date or timestamp.",
arg1
));
// TODO: add once ANSI support is added:
// "When both of the input parameters are not NULL and day_of_week is an invalid input, the function throws SparkIllegalArgumentException if spark.sql.ansi.enabled is set to true, otherwise NULL."
export_functions!((
next_day,
"Returns the first date which is later than start_date and named as indicated. The function returns NULL if at least one of the input parameters is NULL.",
"Returns the first date which is later than start_date and named as indicated. The function returns NULL if at least one of the input parameters is NULL. When both input parameters are not NULL and day_of_week is an invalid input, the function raises an error if `datafusion.execution.enable_ansi_mode` is set to true, otherwise it returns NULL.",
arg1 arg2
));
export_functions!((
Expand Down
312 changes: 188 additions & 124 deletions datafusion/spark/src/function/datetime/next_day.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

use std::sync::Arc;

use arrow::array::{ArrayRef, AsArray, Date32Array, StringArrayType};
use arrow::array::{Array, ArrayRef, AsArray, Date32Array, StringArrayType};
use arrow::datatypes::{DataType, Date32Type, Field, FieldRef};
use chrono::{Datelike, Weekday};
use datafusion_common::{Result, ScalarValue, exec_err, internal_err};
Expand Down Expand Up @@ -69,7 +69,12 @@ impl ScalarUDFImpl for SparkNextDay {
}

fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
let ScalarFunctionArgs { args, .. } = args;
let ScalarFunctionArgs {
args,
config_options,
..
} = args;
let ansi_mode = config_options.execution.enable_ansi_mode;
let [date, day_of_week] = args.as_slice() else {
return exec_err!(
"Spark `next_day` function requires 2 arguments, got {}",
Expand All @@ -86,15 +91,21 @@ impl ScalarUDFImpl for SparkNextDay {
| ScalarValue::LargeUtf8(day_of_week)
| ScalarValue::Utf8View(day_of_week),
) => {
if let Some(days) = days {
if let Some(day_of_week) = day_of_week {
Ok(ColumnarValue::Scalar(ScalarValue::Date32(
spark_next_day(*days, day_of_week.as_str()),
)))
} else {
// TODO: if spark.sql.ansi.enabled is false,
// returns NULL instead of an error for a malformed dayOfWeek.
Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
// Spark's `NextDay` is null intolerant, so a NULL in either
// argument short circuits to NULL even under ANSI mode.
if let (Some(days), Some(day_of_week)) = (days, day_of_week) {
match parse_day_of_week(day_of_week.as_str()) {
Some(weekday) => {
Ok(ColumnarValue::Scalar(ScalarValue::Date32(
next_date_for_day_of_week(*days, weekday),
)))
}
None if ansi_mode => {
illegal_day_of_week_err(day_of_week.as_str())
}
None => {
Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
}
}
} else {
Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
Expand All @@ -113,142 +124,153 @@ impl ScalarUDFImpl for SparkNextDay {
| ScalarValue::LargeUtf8(day_of_week)
| ScalarValue::Utf8View(day_of_week),
) => {
if let Some(day_of_week) = day_of_week {
let result: Date32Array = date_array
.as_primitive::<Date32Type>()
.unary_opt(|days| {
spark_next_day(days, day_of_week.as_str())
})
.with_data_type(DataType::Date32);
Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
} else {
// TODO: if spark.sql.ansi.enabled is false,
// returns NULL instead of an error for a malformed dayOfWeek.
Ok(ColumnarValue::Scalar(ScalarValue::Date32(None)))
match day_of_week
.as_ref()
.map(|d| (d.as_str(), parse_day_of_week(d.as_str())))
{
Some((_, Some(weekday))) => {
let result: Date32Array = date_array
.as_primitive::<Date32Type>()
.unary_opt(|days| {
next_date_for_day_of_week(days, weekday)
})
.with_data_type(DataType::Date32);
Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
}
// The day name is unparsable. Spark's `NextDay` is null
// intolerant per row, so a NULL start date short circuits
// to NULL before the day name is validated. Raise only if
// at least one row has a non-NULL start date, which is
// exactly when the row-wise `process_next_day_arrays` path
// raises for the same inputs.
Some((raw, None))
if ansi_mode
&& date_array.null_count() < date_array.len() =>
{
illegal_day_of_week_err(raw)
}
// Every remaining case yields NULL for every row: an
// unparsable day name with ANSI mode off or with no
// non-NULL start date, or a NULL `day_of_week`.
_ => Ok(ColumnarValue::Scalar(ScalarValue::Date32(None))),
}
}
_ => exec_err!(
"Spark `next_day` function: first arg must be date, second arg must be string. Got {args:?}"
),
}
}
(ColumnarValue::Scalar(date), ColumnarValue::Array(day_of_week_array)) => {
let date_array = date.to_array_of_size(day_of_week_array.len())?;
Ok(ColumnarValue::Array(next_day_arrays(
&date_array,
day_of_week_array,
ansi_mode,
)?))
}
(
ColumnarValue::Array(date_array),
ColumnarValue::Array(day_of_week_array),
) => {
let result = match (date_array.data_type(), day_of_week_array.data_type())
{
(
DataType::Date32,
DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View,
) => {
let date_array: &Date32Array =
date_array.as_primitive::<Date32Type>();
match day_of_week_array.data_type() {
DataType::Utf8 => {
let day_of_week_array =
day_of_week_array.as_string::<i32>();
process_next_day_arrays(date_array, day_of_week_array)
}
DataType::LargeUtf8 => {
let day_of_week_array =
day_of_week_array.as_string::<i64>();
process_next_day_arrays(date_array, day_of_week_array)
}
DataType::Utf8View => {
let day_of_week_array =
day_of_week_array.as_string_view();
process_next_day_arrays(date_array, day_of_week_array)
}
other => {
exec_err!(
"Spark `next_day` function: second arg must be string. Got {other:?}"
)
}
}
}
(left, right) => {
exec_err!(
"Spark `next_day` function: first arg must be date, second arg must be string. Got {left:?}, {right:?}"
)
}
}?;
Ok(ColumnarValue::Array(result))
}
_ => exec_err!("Unsupported args {args:?} for Spark function `next_day`"),
) => Ok(ColumnarValue::Array(next_day_arrays(
date_array,
day_of_week_array,
ansi_mode,
)?)),
}
}
}

fn next_day_arrays(
date_array: &ArrayRef,
day_of_week_array: &ArrayRef,
ansi_mode: bool,
) -> Result<ArrayRef> {
match (date_array.data_type(), day_of_week_array.data_type()) {
(DataType::Date32, DataType::Utf8) => process_next_day_arrays(
date_array.as_primitive::<Date32Type>(),
day_of_week_array.as_string::<i32>(),
ansi_mode,
),
(DataType::Date32, DataType::LargeUtf8) => process_next_day_arrays(
date_array.as_primitive::<Date32Type>(),
day_of_week_array.as_string::<i64>(),
ansi_mode,
),
(DataType::Date32, DataType::Utf8View) => process_next_day_arrays(
date_array.as_primitive::<Date32Type>(),
day_of_week_array.as_string_view(),
ansi_mode,
),
(left, right) => exec_err!(
"Spark `next_day` function: first arg must be date, second arg must be string. Got {left:?}, {right:?}"
),
}
}

fn process_next_day_arrays<'a, S>(
date_array: &Date32Array,
day_of_week_array: &'a S,
ansi_mode: bool,
) -> Result<ArrayRef>
where
&'a S: StringArrayType<'a>,
{
let result = date_array
.iter()
.zip(day_of_week_array.iter())
.map(|(days, day_of_week)| {
if let Some(days) = days {
if let Some(day_of_week) = day_of_week {
spark_next_day(days, day_of_week)
} else {
// TODO: if spark.sql.ansi.enabled is false,
// returns NULL instead of an error for a malformed dayOfWeek.
None
}
} else {
None
let mut builder = Date32Array::builder(date_array.len());
for (days, day_of_week) in date_array.iter().zip(day_of_week_array.iter()) {
// Spark's `NextDay` is null intolerant, so a NULL in either argument
// short circuits to NULL even under ANSI mode.
let (Some(days), Some(day_of_week)) = (days, day_of_week) else {
builder.append_null();
continue;
};
match parse_day_of_week(day_of_week) {
Some(weekday) => {
builder.append_option(next_date_for_day_of_week(days, weekday))
}
})
.collect::<Date32Array>();
Ok(Arc::new(result) as ArrayRef)
None if ansi_mode => return illegal_day_of_week_err(day_of_week),
None => builder.append_null(),
}
}
Ok(Arc::new(builder.finish()) as ArrayRef)
}

fn spark_next_day(days: i32, day_of_week: &str) -> Option<i32> {
let date = Date32Type::to_naive_date_opt(days)?;
/// Mirrors Spark's `DateTimeUtils.getDayOfWeekFromString`: case insensitive, and
/// with no whitespace trimming.
fn parse_day_of_week(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,
}
}

let day_of_week = day_of_week.to_uppercase();
let day_of_week = match day_of_week.as_str() {
"MO" | "MON" | "MONDAY" => Some("MONDAY"),
"TU" | "TUE" | "TUESDAY" => Some("TUESDAY"),
"WE" | "WED" | "WEDNESDAY" => Some("WEDNESDAY"),
"TH" | "THU" | "THURSDAY" => Some("THURSDAY"),
"FR" | "FRI" | "FRIDAY" => Some("FRIDAY"),
"SA" | "SAT" | "SATURDAY" => Some("SATURDAY"),
"SU" | "SUN" | "SUNDAY" => Some("SUNDAY"),
_ => {
// TODO: if spark.sql.ansi.enabled is false,
// returns NULL instead of an error for a malformed dayOfWeek.
None
}
};
/// The first date strictly after `days` that falls on `weekday`. Equivalent to
/// Spark's `DateTimeUtils.getNextDateForDayOfWeek`, so a start date already on
/// `weekday` advances a full week.
///
/// Computes the result on the epoch day directly instead of constructing a
/// `NaiveDate` for it: the result can land past `NaiveDate::MAX` (epoch day
/// 95026236), and building that date panics (`NaiveDate + TimeDelta
/// overflowed`). Spark's `getNextDateForDayOfWeek` is pure `Int` arithmetic and
/// keeps producing a value up to `Int.MaxValue`.
///
/// Returns `None` when `days` is not a representable date, since the weekday is
/// still derived via `NaiveDate`.
fn next_date_for_day_of_week(days: i32, weekday: Weekday) -> Option<i32> {
let date = Date32Type::to_naive_date_opt(days)?;
let delta = 7 - date.weekday().days_since(weekday) as i32;
days.checked_add(delta)
}

if let Some(day_of_week) = day_of_week {
let day_of_week = day_of_week.parse::<Weekday>();
match day_of_week {
Ok(day_of_week) => {
// Advance 1..=7 days from `days` to the next occurrence of
// `day_of_week`. Compute the result on the epoch day directly
// instead of constructing a `NaiveDate`: the result can land
// past `NaiveDate::MAX` (epoch day 95026236), and building that
// date panics (`NaiveDate + TimeDelta overflowed`). Spark's
// `DateTimeUtils.getNextDateForDayOfWeek` is pure `Int`
// arithmetic and keeps producing a value up to `Int.MaxValue`.
let delta = 7 - date.weekday().days_since(day_of_week) as i32;
days.checked_add(delta)
}
Err(_) => {
// TODO: if spark.sql.ansi.enabled is false,
// returns NULL instead of an error for a malformed dayOfWeek.
None
}
}
} else {
None
}
/// Matches Spark's `ILLEGAL_DAY_OF_WEEK` error, raised when ANSI mode is enabled
/// and `day_of_week` does not name a day.
fn illegal_day_of_week_err<T>(day_of_week: &str) -> Result<T> {
exec_err!("Illegal input for day of week: {day_of_week}.")
}

#[cfg(test)]
Expand Down Expand Up @@ -288,8 +310,44 @@ mod tests {

#[test]
fn next_day_rejects_whitespace_padded_day_names() {
assert_eq!(parse_day_of_week(" MO "), None);
assert_eq!(parse_day_of_week("MO "), None);
assert_eq!(parse_day_of_week(""), None);
assert_eq!(parse_day_of_week("mo"), Some(Weekday::Mon));
}

#[test]
fn next_day_advances_a_full_week_on_the_same_weekday() {
let monday = 19723; // 2024-01-01
assert_eq!(spark_next_day(monday, " MO "), None);
assert_eq!(
next_date_for_day_of_week(monday, Weekday::Mon),
Some(monday + 7)
);
assert_eq!(
next_date_for_day_of_week(monday, Weekday::Tue),
Some(monday + 1)
);
}

#[test]
fn next_day_returns_values_past_the_last_representable_date() {
// `chrono::NaiveDate::MAX` is +262142-12-31, which is a Monday. The
// result is computed on the epoch day, so a next occurrence past that
// date still produces a value, as Spark's `Int` arithmetic does.
let max = 95026236;
assert_eq!(next_date_for_day_of_week(max, Weekday::Mon), Some(max + 7));
assert_eq!(next_date_for_day_of_week(max - 1, Weekday::Mon), Some(max));
// A start date already on the requested weekday advances a full week,
// so it lands one day past `max` from six days earlier.
assert_eq!(
next_date_for_day_of_week(max - 6, Weekday::Tue),
Some(max + 1)
);
assert_eq!(next_date_for_day_of_week(max - 6, Weekday::Mon), Some(max));
// The weekday is still derived via `NaiveDate`, so a *start* day that is
// not a representable date returns NULL.
assert_eq!(next_date_for_day_of_week(i32::MAX, Weekday::Mon), None);
assert_eq!(next_date_for_day_of_week(i32::MIN, Weekday::Mon), None);
}

#[test]
Expand All @@ -301,7 +359,13 @@ mod tests {
// than panicking with `NaiveDate + TimeDelta overflowed`.
//
// 95026236 is a Monday, so `next_day(.., "Mon")` advances a full week.
assert_eq!(spark_next_day(95026236, "Mon"), Some(95026243));
assert_eq!(spark_next_day(95026230, "Tue"), Some(95026237));
assert_eq!(
next_date_for_day_of_week(95026236, Weekday::Mon),
Some(95026243)
);
assert_eq!(
next_date_for_day_of_week(95026230, Weekday::Tue),
Some(95026237)
);
}
}
Loading