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 1 commit into
base: main
Choose a base branch
from
+503 −135
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
1 commit
Select commit
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
288 changes: 172 additions & 116 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, Duration, 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,136 +124,151 @@ 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.
///
/// Returns `None` when `days` is not a representable date, and also when the
/// result would fall past `chrono::NaiveDate::MAX`. Spark computes the answer
/// with plain `Int` arithmetic and has no calendar range limit, so it keeps
/// returning values in that far-future range where this returns NULL.
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 i64;
Some(Date32Type::from_naive_date(
date.checked_add_signed(Duration::days(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) => Some(Date32Type::from_naive_date(
date + Duration::days(
(7 - date.weekday().days_since(day_of_week)) as i64,
),
)),
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 @@ -282,7 +308,37 @@ 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_null_past_the_last_representable_date() {
// `chrono::NaiveDate::MAX` is +262142-12-31, which is a Monday.
let max = 95026236;
assert_eq!(next_date_for_day_of_week(max, Weekday::Mon), None);
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 overflows from six days earlier.
assert_eq!(next_date_for_day_of_week(max - 6, Weekday::Tue), None);
assert_eq!(next_date_for_day_of_week(max - 6, Weekday::Mon), Some(max));
// Values that are not representable dates at all still return 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);
}
}
Loading
Loading