From 37768b8aadd0e9de543029778814d585588ff64f Mon Sep 17 00:00:00 2001 From: pchintar <89355405+pchintar@users.noreply.github.com> Date: Sat, 13 Jun 2026 07:04:43 +0530 Subject: [PATCH 1/4] fix: Enable sliding window execution for covar_pop, covar_samp, and corr (#22764) ## Which issue does this PR close? - Closes #22763 . ## Rationale for this change Bounded sliding window queries using `covar_pop`, `covar_samp`, and `corr` currently fail with a `retract_batch is not implemented` error, preventing these aggregates from being used with sliding window frames. ## What changes are included in this PR? * Included `supports_retract_batch()` for the covariance and correlation accumulators. * Added SQL logic tests covering bounded sliding window execution for covariance and correlation aggregates. ## Are these changes tested? Yes. Added SQL logic tests covering: * Single-row bounded sliding frames * Multi-row bounded sliding frames for `covar_pop`, `covar_samp`, and `corr`. ## Are there any user-facing changes? Yes. `covar_pop`, `covar_samp`, and `corr` can now be used with bounded sliding window frames that previously failed. Also, no changes were made to any public APIs. --- .../functions-aggregate/src/correlation.rs | 4 + .../functions-aggregate/src/covariance.rs | 12 ++ .../functions-aggregate/src/variance.rs | 7 + datafusion/sqllogictest/test_files/window.slt | 136 ++++++++++++++++++ 4 files changed, 159 insertions(+) diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index 2621fcf0bf3c7..5a95cfe8320fc 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -281,6 +281,10 @@ impl Accumulator for CorrelationAccumulator { self.stddev2.retract_batch(&values[1..2])?; Ok(()) } + + fn supports_retract_batch(&self) -> bool { + true + } } #[derive(Default)] diff --git a/datafusion/functions-aggregate/src/covariance.rs b/datafusion/functions-aggregate/src/covariance.rs index 18d602ab33940..bd7c8a039076a 100644 --- a/datafusion/functions-aggregate/src/covariance.rs +++ b/datafusion/functions-aggregate/src/covariance.rs @@ -305,6 +305,14 @@ impl Accumulator for CovarianceAccumulator { _ => continue, }; + if self.count <= 1 { + self.count = 0; + self.mean1 = 0.0; + self.mean2 = 0.0; + self.algo_const = 0.0; + continue; + } + let new_count = self.count - 1; let delta1 = self.mean1 - value1; let new_mean1 = delta1 / new_count as f64 + self.mean1; @@ -373,4 +381,8 @@ impl Accumulator for CovarianceAccumulator { fn size(&self) -> usize { size_of_val(self) } + + fn supports_retract_batch(&self) -> bool { + true + } } diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index ce3e00b9ffd91..d5fddf01f2d52 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -348,6 +348,13 @@ impl Accumulator for VarianceAccumulator { fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let arr = as_float64_array(&values[0])?; for value in arr.iter().flatten() { + if self.count <= 1 { + self.count = 0; + self.mean = 0.0; + self.m2 = 0.0; + continue; + } + let new_count = self.count - 1; let delta1 = self.mean - value; let new_mean = delta1 / new_count as f64 + self.mean; diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 1b51950a70e1b..59b43d6476571 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -6626,6 +6626,142 @@ ORDER BY i; 3 1 4 NULL +# Covariance/correlation sliding-window regression test. Verifies correct +# results across row removals and a NULL-gap empty-frame transition. +query IRRR +SELECT + column1, + covar_pop(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + covar_samp(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + corr(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0, 5.0), + (2, NULL, NULL), + (3, NULL, NULL), + (4, 30.0, 10.0), + (5, 40.0, 20.0), + (6, 50.0, 10.0) +); +---- +1 0 NULL NULL +2 0 NULL NULL +3 NULL NULL NULL +4 0 NULL NULL +5 25 50 1 +6 -25 -50 -1 + +# Multi-row covariance/correlation sliding-window regression test. Verifies +# correct accumulation when valid rows enter the frame after a reset. +query IRRR +SELECT + column1, + covar_pop(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 2 PRECEDING AND CURRENT ROW + ), + covar_samp(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 2 PRECEDING AND CURRENT ROW + ), + corr(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 2 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0, 5.0), + (2, NULL, NULL), + (3, NULL, NULL), + (4, 30.0, 10.0), + (5, 40.0, 20.0), + (6, 50.0, 10.0) +); +---- +1 0 NULL NULL +2 0 NULL NULL +3 0 NULL NULL +4 0 NULL NULL +5 25 50 1 +6 0 0 0 + +# Covariance/correlation sliding-window regression test. Rows with NULL in +# either input column must not contribute to the aggregate state. +query IRRR +SELECT + column1, + covar_pop(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 3 PRECEDING AND CURRENT ROW + ), + covar_samp(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 3 PRECEDING AND CURRENT ROW + ), + corr(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 3 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0, 5.0), + (2, 20.0, NULL), + (3, NULL, 15.0), + (4, 30.0, 10.0), + (5, 40.0, 20.0) +); +---- +1 0 NULL NULL +2 0 NULL NULL +3 0 NULL NULL +4 25 50 1 +5 25 50 1 + +# Variance/stddev sliding-window regression test. Verifies that retracting +# the last valid row resets the aggregate state. +query IRRRR +SELECT + column1, + var_pop(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + var_samp(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + stddev_pop(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + stddev_samp(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0), + (2, NULL), + (3, NULL), + (4, 30.0), + (5, 40.0) +); +---- +1 0 NULL 0 NULL +2 0 NULL 0 NULL +3 NULL NULL NULL NULL +4 0 NULL 0 NULL +5 25 50 5 7.071067811865 + # Decimal variant — the integer-division path would otherwise panic on an # empty frame. query IR From f9317284cf56c517742c212fe893af05ff955a23 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Sat, 13 Jun 2026 07:06:39 +0530 Subject: [PATCH 2/4] fix: handle `date_bin` negative subsecond and overflow cases (#22610) ## Which issue does this PR close? - Closes #22528 ## Rationale for this change `date_bin` had a few edge cases that could return the wrong result, return an error only on array inputs, or panic/wrap when scaling timestamp and time values to nanoseconds. ## What changes are included in this PR? - Fix negative sub-second timestamp conversion before the epoch. - Make scalar and array paths return `NULL` consistently for per-row binning errors. - Use checked scaling when converting timestamp and time values to nanoseconds. - Return an error for invalid shared origin values that overflow during scaling. - Simplify duplicated stride and scale handling. ## Are these changes tested? Yes ## Are there any user-facing changes? No public API changes. --- datafusion/functions/src/datetime/date_bin.rs | 439 +++++++++--------- .../test_files/date_bin_errors.slt | 28 +- 2 files changed, 250 insertions(+), 217 deletions(-) diff --git a/datafusion/functions/src/datetime/date_bin.rs b/datafusion/functions/src/datetime/date_bin.rs index 38b491e42bcbd..06ffd8ba5b3c6 100644 --- a/datafusion/functions/src/datetime/date_bin.rs +++ b/datafusion/functions/src/datetime/date_bin.rs @@ -34,9 +34,7 @@ use arrow::datatypes::{ use arrow::error::ArrowError; use arrow::temporal_conversions::NANOSECONDS_IN_DAY; use datafusion_common::cast::as_primitive_array; -use datafusion_common::{ - DataFusionError, Result, ScalarValue, exec_err, not_impl_err, plan_err, -}; +use datafusion_common::{Result, ScalarValue, exec_err, not_impl_err, plan_err}; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ @@ -420,14 +418,44 @@ fn date_bin_months_interval(stride_months: i64, source: i64, origin: i64) -> Res } fn to_utc_date_time(nanos: i64) -> Result> { - let secs = nanos / NANOS_PER_SEC; - let nsec = (nanos % NANOS_PER_SEC) as u32; + // Keep negative sub-second values normalized as seconds + non-negative nanos. + let secs = nanos.div_euclid(NANOS_PER_SEC); + let nsec = nanos.rem_euclid(NANOS_PER_SEC) as u32; match DateTime::from_timestamp(secs, nsec) { Some(dt) => Ok(dt), None => exec_err!("Invalid timestamp value"), } } +fn timestamp_scale() -> i64 { + match T::UNIT { + Nanosecond => 1, + Microsecond => NANOS_PER_MICRO, + Millisecond => NANOS_PER_MILLI, + Second => NANOSECONDS, + } +} + +// Scale to nanoseconds and report overflow as a normal error. +fn checked_scale_to_nanos(x: i64, scale: i64) -> Result { + match x.checked_mul(scale) { + Some(scaled) => Ok(scaled), + None => exec_err!("date_bin timestamp value {x} * scale {scale} overflows i64"), + } +} + +fn validate_time_stride(stride: &Interval) -> Result<()> { + match stride { + Interval::Months(m) if *m > 0 => { + exec_err!("DATE_BIN stride for TIME input must be less than 1 day") + } + Interval::Nanoseconds(ns) if *ns >= NANOSECONDS_IN_DAY => { + exec_err!("DATE_BIN stride for TIME input must be less than 1 day") + } + _ => Ok(()), + } +} + // Supported intervals: // 1. IntervalDayTime: this means that the stride is in days, hours, minutes, seconds and milliseconds // We will assume month interval won't be converted into this type @@ -498,83 +526,20 @@ fn date_bin_impl( (*v, false) } ColumnarValue::Scalar(ScalarValue::Time32Millisecond(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - - (*v as i64 * NANOS_PER_MILLI, true) + validate_time_stride(&stride)?; + // TIME origins can come from reinterpret casts, so scale defensively. + (checked_scale_to_nanos(*v as i64, NANOS_PER_MILLI)?, true) } ColumnarValue::Scalar(ScalarValue::Time32Second(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - - (*v as i64 * NANOS_PER_SEC, true) + validate_time_stride(&stride)?; + (checked_scale_to_nanos(*v as i64, NANOS_PER_SEC)?, true) } ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - - (*v * NANOS_PER_MICRO, true) + validate_time_stride(&stride)?; + (checked_scale_to_nanos(*v, NANOS_PER_MICRO)?, true) } ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - + validate_time_stride(&stride)?; (*v, true) } ColumnarValue::Scalar(v) => { @@ -597,91 +562,49 @@ fn date_bin_impl( return exec_err!("DATE_BIN stride must be non-zero"); } - fn timestamp_scale() -> i64 { - match T::UNIT { - Nanosecond => 1, - Microsecond => NANOS_PER_MICRO, - Millisecond => NANOS_PER_MILLI, - Second => NANOSECONDS, - } - } - - fn timestamp_scale_overflow_error(x: i64) -> DataFusionError { - DataFusionError::Execution(format!( - "DATE_BIN source timestamp {x} cannot be represented in nanoseconds" - )) + fn transform_scalar_with_stride( + value: Option, + origin: i64, + stride: i64, + stride_fn: BinFunction, + ) -> Option { + let scale = timestamp_scale::(); + value + .and_then(|val| val.checked_mul(scale)) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| binned / scale) } Ok(match array { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(v, tz_opt)) => { - let scale = timestamp_scale::(); ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( - match *v { - Some(val) => { - let scaled = val - .checked_mul(scale) - .ok_or_else(|| timestamp_scale_overflow_error(val))?; - match stride_fn(stride, scaled, origin) { - Ok(result) => Some(result / scale), - Err(_) => None, - } - } - None => None, - }, + transform_scalar_with_stride::( + *v, origin, stride, stride_fn, + ), tz_opt.clone(), )) } ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(v, tz_opt)) => { - let scale = timestamp_scale::(); ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond( - match *v { - Some(val) => { - let scaled = val - .checked_mul(scale) - .ok_or_else(|| timestamp_scale_overflow_error(val))?; - match stride_fn(stride, scaled, origin) { - Ok(result) => Some(result / scale), - Err(_) => None, - } - } - None => None, - }, + transform_scalar_with_stride::( + *v, origin, stride, stride_fn, + ), tz_opt.clone(), )) } ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(v, tz_opt)) => { - let scale = timestamp_scale::(); ColumnarValue::Scalar(ScalarValue::TimestampMillisecond( - match *v { - Some(val) => { - let scaled = val - .checked_mul(scale) - .ok_or_else(|| timestamp_scale_overflow_error(val))?; - match stride_fn(stride, scaled, origin) { - Ok(result) => Some(result / scale), - Err(_) => None, - } - } - None => None, - }, + transform_scalar_with_stride::( + *v, origin, stride, stride_fn, + ), tz_opt.clone(), )) } ColumnarValue::Scalar(ScalarValue::TimestampSecond(v, tz_opt)) => { - let scale = timestamp_scale::(); ColumnarValue::Scalar(ScalarValue::TimestampSecond( - match *v { - Some(val) => { - let scaled = val - .checked_mul(scale) - .ok_or_else(|| timestamp_scale_overflow_error(val))?; - match stride_fn(stride, scaled, origin) { - Ok(result) => Some(result / scale), - Err(_) => None, - } - } - None => None, - }, + transform_scalar_with_stride::( + *v, origin, stride, stride_fn, + ), tz_opt.clone(), )) } @@ -689,39 +612,30 @@ fn date_bin_impl( if !is_time { return exec_err!("DATE_BIN with Time32 source requires Time32 origin"); } - let result = v.and_then(|x| { - match stride_fn(stride, x as i64 * NANOS_PER_MILLI, origin) { - Ok(binned_nanos) => { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - Some((nanos / NANOS_PER_MILLI) as i32) - } - Err(_) => None, - } - }); + let result = v + .and_then(|x| (x as i64).checked_mul(NANOS_PER_MILLI)) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| ((binned % NANOSECONDS_IN_DAY) / NANOS_PER_MILLI) as i32); ColumnarValue::Scalar(ScalarValue::Time32Millisecond(result)) } ColumnarValue::Scalar(ScalarValue::Time32Second(v)) => { if !is_time { return exec_err!("DATE_BIN with Time32 source requires Time32 origin"); } - let result = v.and_then(|x| { - match stride_fn(stride, x as i64 * NANOS_PER_SEC, origin) { - Ok(binned_nanos) => { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - Some((nanos / NANOS_PER_SEC) as i32) - } - Err(_) => None, - } - }); + let result = v + .and_then(|x| (x as i64).checked_mul(NANOS_PER_SEC)) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| ((binned % NANOSECONDS_IN_DAY) / NANOS_PER_SEC) as i32); ColumnarValue::Scalar(ScalarValue::Time32Second(result)) } ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(v)) => { if !is_time { return exec_err!("DATE_BIN with Time64 source requires Time64 origin"); } - let result = v.and_then(|x| match stride_fn(stride, x, origin) { - Ok(binned_nanos) => Some(binned_nanos % (NANOSECONDS_IN_DAY)), - Err(_) => None, + let result = v.and_then(|x| { + stride_fn(stride, x, origin) + .map(|binned| binned % NANOSECONDS_IN_DAY) + .ok() }); ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(result)) } @@ -729,14 +643,10 @@ fn date_bin_impl( if !is_time { return exec_err!("DATE_BIN with Time64 source requires Time64 origin"); } - let result = - v.and_then(|x| match stride_fn(stride, x * NANOS_PER_MICRO, origin) { - Ok(binned_nanos) => { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - Some(nanos / NANOS_PER_MICRO) - } - Err(_) => None, - }); + let result = v + .and_then(|x| x.checked_mul(NANOS_PER_MICRO)) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| (binned % NANOSECONDS_IN_DAY) / NANOS_PER_MICRO); ColumnarValue::Scalar(ScalarValue::Time64Microsecond(result)) } ColumnarValue::Array(array) => { @@ -753,22 +663,12 @@ fn date_bin_impl( let array = as_primitive_array::(array)?; let scale = timestamp_scale::(); - let values = array - .iter() - .map(|val| match val { - Some(val) => { - let scaled = val - .checked_mul(scale) - .ok_or_else(|| timestamp_scale_overflow_error(val))?; - Ok(stride_fn(stride, scaled, origin) - .ok() - .map(|binned| binned / scale)) - } - None => Ok(None), - }) - .collect::>>()?; - - let result = PrimitiveArray::::from_iter(values); + // Per-row errors become NULL, matching scalar behavior. + let result: PrimitiveArray = array.unary_opt(|val| { + val.checked_mul(scale) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| binned / scale) + }); let array = result.with_timezone_opt(tz_opt.clone()); Ok(ColumnarValue::Array(Arc::new(array))) @@ -803,14 +703,15 @@ fn date_bin_impl( } let array = array.as_primitive::(); let result: PrimitiveArray = - array.try_unary(|x| { - stride_fn(stride, x as i64 * NANOS_PER_MILLI, origin) - .map(|binned_nanos| { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - (nanos / NANOS_PER_MILLI) as i32 + array.unary_opt(|x| { + (x as i64) + .checked_mul(NANOS_PER_MILLI) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| { + ((binned % NANOSECONDS_IN_DAY) / NANOS_PER_MILLI) + as i32 }) - .map_err(|e| ArrowError::ComputeError(e.to_string())) - })?; + }); ColumnarValue::Array(Arc::new(result)) } Time32(Second) => { @@ -820,15 +721,14 @@ fn date_bin_impl( ); } let array = array.as_primitive::(); - let result: PrimitiveArray = - array.try_unary(|x| { - stride_fn(stride, x as i64 * NANOS_PER_SEC, origin) - .map(|binned_nanos| { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - (nanos / NANOS_PER_SEC) as i32 - }) - .map_err(|e| ArrowError::ComputeError(e.to_string())) - })?; + let result: PrimitiveArray = array.unary_opt(|x| { + (x as i64) + .checked_mul(NANOS_PER_SEC) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| { + ((binned % NANOSECONDS_IN_DAY) / NANOS_PER_SEC) as i32 + }) + }); ColumnarValue::Array(Arc::new(result)) } Time64(Microsecond) => { @@ -839,14 +739,13 @@ fn date_bin_impl( } let array = array.as_primitive::(); let result: PrimitiveArray = - array.try_unary(|x| { - stride_fn(stride, x * NANOS_PER_MICRO, origin) - .map(|binned_nanos| { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - nanos / NANOS_PER_MICRO + array.unary_opt(|x| { + x.checked_mul(NANOS_PER_MICRO) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) + .map(|binned| { + (binned % NANOSECONDS_IN_DAY) / NANOS_PER_MICRO }) - .map_err(|e| ArrowError::ComputeError(e.to_string())) - })?; + }); ColumnarValue::Array(Arc::new(result)) } Time64(Nanosecond) => { @@ -857,11 +756,11 @@ fn date_bin_impl( } let array = array.as_primitive::(); let result: PrimitiveArray = - array.try_unary(|x| { + array.unary_opt(|x| { stride_fn(stride, x, origin) .map(|binned_nanos| binned_nanos % (NANOSECONDS_IN_DAY)) - .map_err(|e| ArrowError::ComputeError(e.to_string())) - })?; + .ok() + }); ColumnarValue::Array(Arc::new(result)) } _ => { @@ -917,6 +816,31 @@ mod tests { DateBinFunc::new().invoke_with_args(args) } + fn assert_null_scalar(value: ColumnarValue, expected_type: DataType) { + let ColumnarValue::Scalar(value) = value else { + panic!("expected scalar, got {value:?}"); + }; + assert_eq!(value.data_type(), expected_type); + assert!(value.is_null(), "expected NULL, got {value:?}"); + } + + fn assert_array_null_then_valid(value: ColumnarValue, expected_type: DataType) { + let ColumnarValue::Array(array) = value else { + panic!("expected array, got {value:?}"); + }; + assert_eq!(array.data_type(), &expected_type); + assert!(array.is_null(0), "expected NULL at row 0"); + assert!(array.is_valid(1), "expected valid value at row 1"); + } + + fn assert_overflow_error(result: Result) { + let err = result.expect_err("expected overflow error"); + assert!( + err.strip_backtrace().contains("overflows i64"), + "unexpected error: {err}" + ); + } + #[test] fn test_date_bin() { let return_field = &Arc::new(Field::new( @@ -1433,6 +1357,97 @@ mod tests { } } + #[test] + fn test_date_bin_scale_overflow_returns_null() { + // Scaling non-nanosecond timestamps to nanoseconds can overflow. + use arrow::array::{ + ArrayRef, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampSecondArray, + }; + + let scalar_cases = [ + ScalarValue::TimestampSecond(Some(i64::MAX), None), + ScalarValue::TimestampMillisecond(Some(i64::MAX), None), + ScalarValue::TimestampMicrosecond(Some(i64::MAX), None), + ]; + for source in scalar_cases { + let expected_type = source.data_type(); + let return_field = Arc::new(Field::new("f", expected_type.clone(), true)); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), + ColumnarValue::Scalar(source), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)), + ]; + let result = invoke_date_bin_with_args(args, 1, &return_field) + .unwrap_or_else(|e| panic!("expected Ok for {expected_type}, got {e:?}")); + assert_null_scalar(result, expected_type); + } + + let array_cases: Vec = vec![ + Arc::new(TimestampSecondArray::from(vec![Some(i64::MAX), Some(0)])), + Arc::new(TimestampMillisecondArray::from(vec![ + Some(i64::MAX), + Some(0), + ])), + Arc::new(TimestampMicrosecondArray::from(vec![ + Some(i64::MAX), + Some(0), + ])), + ]; + for array in array_cases { + let dt = array.data_type().clone(); + let return_field = Arc::new(Field::new("f", dt.clone(), true)); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), + ColumnarValue::Array(array), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)), + ]; + let result = invoke_date_bin_with_args(args, 2, &return_field) + .unwrap_or_else(|e| panic!("expected Ok for {dt:?}, got {e:?}")); + assert_array_null_then_valid(result, dt); + } + } + + #[test] + fn test_date_bin_time64_micro_overflow_handling() { + // Time64(Microsecond) can hold out-of-range values after reinterpret casts. + use arrow::array::Time64MicrosecondArray; + + let data_type = DataType::Time64(TimeUnit::Microsecond); + let return_field = &Arc::new(Field::new("f", data_type.clone(), true)); + let stride = || ColumnarValue::Scalar(ScalarValue::new_interval_dt(0, 1000)); + let origin = || ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(0))); + + // Out-of-range source values are per-row data, so they become NULL. + let args = vec![ + stride(), + ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(i64::MAX))), + origin(), + ]; + let result = invoke_date_bin_with_args(args, 1, return_field).unwrap(); + assert_null_scalar(result, data_type.clone()); + + let array = Arc::new(Time64MicrosecondArray::from(vec![Some(i64::MAX), Some(0)])); + let args = vec![stride(), ColumnarValue::Array(array), origin()]; + let result = invoke_date_bin_with_args(args, 2, return_field).unwrap(); + assert_array_null_then_valid(result, data_type); + + let bad_origin = + || ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(i64::MAX))); + + // Out-of-range origins are shared inputs, so they return an error. + let args = vec![ + stride(), + ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(0))), + bad_origin(), + ]; + assert_overflow_error(invoke_date_bin_with_args(args, 1, return_field)); + + let array = Arc::new(Time64MicrosecondArray::from(vec![Some(0), Some(1)])); + let args = vec![stride(), ColumnarValue::Array(array), bad_origin()]; + assert_overflow_error(invoke_date_bin_with_args(args, 2, return_field)); + } + #[test] fn test_date_bin_compute_distance_rem_overflow() { // Regression for #22215: `time_diff % stride` panics with "attempt to diff --git a/datafusion/sqllogictest/test_files/date_bin_errors.slt b/datafusion/sqllogictest/test_files/date_bin_errors.slt index 20408c84ef79a..53cba506defd6 100644 --- a/datafusion/sqllogictest/test_files/date_bin_errors.slt +++ b/datafusion/sqllogictest/test_files/date_bin_errors.slt @@ -23,10 +23,24 @@ select date_bin(interval '1637426858 months', to_timestamp_millis(1040292460), t ---- NULL -# Negative timestamp with month interval - should return NULL instead of panicking +# Issue #22528: negative sub-second source with month interval. query P select date_bin(interval '1 month', to_timestamp_millis(-1040292460), timestamp '1984-01-07 00:00:00'); ---- +1969-12-07T00:00:00 + +# Array path should match the scalar path above. +query P +select date_bin(interval '1 month', c, timestamp '1984-01-07 00:00:00') +from values (to_timestamp_millis(-1040292460)) t(c); +---- +1969-12-07T00:00:00 + +# Array path should return NULL for per-row overflow. +query P +select date_bin(interval '1637426858 months', c, timestamp '1984-01-07 00:00:00') +from values (to_timestamp_millis(1040292460)) t(c); +---- NULL # Large stride causing overflow - should return NULL @@ -79,16 +93,18 @@ select date_bin( ---- NULL -# Source timestamp scaling to nanoseconds overflows: should return an error, not panic -query error DataFusion error: Execution error: DATE_BIN source timestamp 9223372036854775807 cannot be represented in nanoseconds +# Source timestamp scaling to nanoseconds overflows: should return NULL, not panic +query P select date_bin( interval '1 nanosecond', arrow_cast(9223372036854775807, 'Timestamp(Second, None)'), timestamp '1970-01-01 00:00:00' ); +---- +NULL -# Source timestamp scaling to nanoseconds overflows in array path: should return an error, not panic -query error DataFusion error: Execution error: DATE_BIN source timestamp 9223372036854775807 cannot be represented in nanoseconds +# Source timestamp scaling to nanoseconds overflows in array path: should return NULL, not panic +query P select date_bin( interval '1 nanosecond', ts, @@ -97,3 +113,5 @@ select date_bin( from ( values (arrow_cast(9223372036854775807, 'Timestamp(Second, None)')) ) as t(ts); +---- +NULL From e5f7af15a6dafe9e2847a81bae5fddfda5a60c0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20L=C3=B3pez?= Date: Sat, 13 Jun 2026 03:37:17 +0200 Subject: [PATCH 3/4] feat(spark): add `concat_ws` with array support (#20928) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Part of #15914 ## Rationale for this change DataFusion core's `concat_ws` does not support array arguments. Spark's `concat_ws(sep, ...)` accepts both scalar strings and arrays, expanding array elements and skipping nulls. This is needed for Spark compatibility in the `datafusion-spark` crate. ## What changes are included in this PR? - New `SparkConcatWs` UDF in `datafusion/spark/src/function/string/concat_ws.rs` - Supports `concat_ws(sep, str1, str2, ...)` with scalar strings - Supports array arguments: `concat_ws(',', array('a', 'b'), 'c')` → `"a,b,c"` - Null scalars and null array elements are skipped (Spark behavior) - Null separator returns NULL - Zero value arguments (`concat_ws(',')`) returns empty string - Supports Utf8, LargeUtf8, Utf8View, List, and LargeList types - Registered the function in `mod.rs` (`make_udf_function!`, `export_functions!`, `functions()`) - Replaced commented-out SLT tests with 14 working test cases covering basic usage, arrays, mixed arguments, nulls, column expressions, and edge cases ## Are these changes tested? Yes. - 7 unit tests in `concat_ws.rs` (basic, null values skipped, null separator, list arrays, list with nulls, mixed scalar+list, multiple rows) - 14 SLT tests in `spark/string/concat_ws.slt` covering scalars, arrays, nulls, column expressions, and edge cases ## Are there any user-facing changes? No. This is a new function in the `datafusion-spark` crate only. --- .../spark/src/function/string/concat_ws.rs | 297 ++++++++++++++ datafusion/spark/src/function/string/mod.rs | 8 + .../test_files/spark/string/concat_ws.slt | 383 +++++++++++++++++- 3 files changed, 669 insertions(+), 19 deletions(-) create mode 100644 datafusion/spark/src/function/string/concat_ws.rs diff --git a/datafusion/spark/src/function/string/concat_ws.rs b/datafusion/spark/src/function/string/concat_ws.rs new file mode 100644 index 0000000000000..c9ed1369a51a7 --- /dev/null +++ b/datafusion/spark/src/function/string/concat_ws.rs @@ -0,0 +1,297 @@ +// 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. + +//! Spark-compatible `concat_ws`: joins strings (and array elements) with a separator. +//! +//! Null scalar args and null array elements are skipped; a null separator yields a +//! null row. Non-string args are coerced to STRING; list args (`List`, `LargeList`, +//! `ListView`, `LargeListView`, `FixedSizeList`) expand their elements. +//! +//! Differences with DataFusion core `concat_ws`: +//! - Accepts list arguments and expands their elements +//! - Always returns Utf8 (Spark's `STRING` type) +//! - Coerces non-string scalars (numbers, booleans, dates, ...) to Utf8 + +use std::fmt::Write as _; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, GenericListArray, LargeStringArray, OffsetSizeTrait, + StringArray, StringBuilder, StringViewArray, +}; +use arrow::datatypes::{DataType, Field}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; + +use crate::function::error_utils::{ + invalid_arg_count_exec_err, unsupported_data_type_exec_err, +}; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkConcatWs { + signature: Signature, +} + +impl Default for SparkConcatWs { + fn default() -> Self { + Self::new() + } +} + +impl SparkConcatWs { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkConcatWs { + fn name(&self) -> &str { + "concat_ws" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.is_empty() { + return Err(invalid_arg_count_exec_err("concat_ws", (1, i32::MAX), 0)); + } + Ok(arg_types + .iter() + .enumerate() + .map(|(i, dt)| match dt { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => dt.clone(), + // Non-separator list args expand their elements at runtime. + // Normalize the list variant so the kernel only sees + // List/LargeList, AND force the element type to Utf8 so the + // planner inserts a cast for non-string children (Spark + // coerces them to STRING the same way it does for scalars). + DataType::List(f) + | DataType::ListView(f) + | DataType::FixedSizeList(f, _) + if i > 0 => + { + DataType::List(Arc::new(Field::new( + f.name(), + DataType::Utf8, + f.is_nullable(), + ))) + } + DataType::LargeList(f) | DataType::LargeListView(f) if i > 0 => { + DataType::LargeList(Arc::new(Field::new( + f.name(), + DataType::Utf8, + f.is_nullable(), + ))) + } + // Spark casts everything else (numbers, booleans, dates, + // binary, null...) to STRING. + _ => DataType::Utf8, + }) + .collect()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + // Only separator provided → empty string (or NULL if separator is null). + // Arg-count validation happens in coerce_types at planning time. + if args.args.len() == 1 { + return only_separator(&args.args[0]); + } + + spark_concat_ws(&args.args, args.number_rows) + } +} + +fn only_separator(sep: &ColumnarValue) -> Result { + match sep { + ColumnarValue::Scalar(s) if s.is_null() => { + Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))) + } + ColumnarValue::Scalar(_) => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some( + String::new(), + )))), + ColumnarValue::Array(arr) => { + let mut builder = StringBuilder::with_capacity(arr.len(), 0); + for row_idx in 0..arr.len() { + if arr.is_null(row_idx) { + builder.append_null(); + } else { + builder.append_value(""); + } + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef)) + } + } +} + +fn spark_concat_ws(args: &[ColumnarValue], num_rows: usize) -> Result { + let arrays = ColumnarValue::values_to_arrays(args)?; + let sep_view = StringView::try_new(&arrays[0])?; + let arg_views: Vec = arrays[1..] + .iter() + .map(ArgView::try_new) + .collect::>()?; + + let mut builder = StringBuilder::with_capacity(num_rows, num_rows * 16); + + for row_idx in 0..num_rows { + if sep_view.is_null(row_idx) { + builder.append_null(); + continue; + } + + // Write parts directly into the builder via its `fmt::Write` impl; + // `append_value("")` then finalises the row (offset + validity) with + // no extra copy from an intermediate `String`. + let separator = sep_view.value(row_idx); + let mut first = true; + for view in &arg_views { + view.write_row(row_idx, separator, &mut builder, &mut first)?; + } + builder.append_value(""); + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef)) +} + +/// Typed view over a string array that downcasts once and exposes +/// per-row access without further dispatch. +enum StringView<'a> { + Utf8(&'a StringArray), + LargeUtf8(&'a LargeStringArray), + Utf8View(&'a StringViewArray), +} + +impl<'a> StringView<'a> { + fn try_new(arr: &'a ArrayRef) -> Result { + match arr.data_type() { + DataType::Utf8 => Ok(Self::Utf8(arr.as_string::())), + DataType::LargeUtf8 => Ok(Self::LargeUtf8(arr.as_string::())), + DataType::Utf8View => Ok(Self::Utf8View(arr.as_string_view())), + other => Err(unsupported_data_type_exec_err("concat_ws", "STRING", other)), + } + } + + fn value(&self, idx: usize) -> &str { + match self { + Self::Utf8(a) => a.value(idx), + Self::LargeUtf8(a) => a.value(idx), + Self::Utf8View(a) => a.value(idx), + } + } + + fn is_null(&self, idx: usize) -> bool { + match self { + Self::Utf8(a) => a.is_null(idx), + Self::LargeUtf8(a) => a.is_null(idx), + Self::Utf8View(a) => a.is_null(idx), + } + } +} + +/// Per-argument view: a string array or a list of strings. The downcast +/// happens once at construction time. `DataType::Null` cannot appear here — +/// `coerce_types` rewrites it to `Utf8` before invocation. +enum ArgView<'a> { + Str(StringView<'a>), + List(&'a GenericListArray), + LargeList(&'a GenericListArray), +} + +impl<'a> ArgView<'a> { + fn try_new(arr: &'a ArrayRef) -> Result { + match arr.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + Ok(Self::Str(StringView::try_new(arr)?)) + } + DataType::List(_) => Ok(Self::List(arr.as_list::())), + DataType::LargeList(_) => Ok(Self::LargeList(arr.as_list::())), + other => Err(unsupported_data_type_exec_err( + "concat_ws", + "STRING or ARRAY", + other, + )), + } + } + + fn write_row( + &self, + row_idx: usize, + sep: &str, + builder: &mut StringBuilder, + first: &mut bool, + ) -> Result<()> { + match self { + Self::Str(view) => { + if !view.is_null(row_idx) { + push_part(builder, view.value(row_idx), sep, first); + } + } + Self::List(list) => write_list_row(*list, row_idx, sep, builder, first)?, + Self::LargeList(list) => write_list_row(*list, row_idx, sep, builder, first)?, + } + Ok(()) + } +} + +fn write_list_row( + list: &GenericListArray, + row_idx: usize, + sep: &str, + builder: &mut StringBuilder, + first: &mut bool, +) -> Result<()> { + if list.is_null(row_idx) { + return Ok(()); + } + let values = list.value(row_idx); + // An empty array (e.g. `array()`) contributes nothing — Spark renders it + // as the empty string, not an error. + if values.is_empty() { + return Ok(()); + } + let view = StringView::try_new(&values)?; + for i in 0..values.len() { + if !view.is_null(i) { + push_part(builder, view.value(i), sep, first); + } + } + Ok(()) +} + +// `StringBuilder::write_str` only does `extend_from_slice` and never errors; +// the `.expect(..)` is a documentation hint, not a real failure path. +fn push_part(builder: &mut StringBuilder, part: &str, sep: &str, first: &mut bool) { + if !*first { + builder + .write_str(sep) + .expect("StringBuilder::write_str is infallible"); + } + *first = false; + builder + .write_str(part) + .expect("StringBuilder::write_str is infallible"); +} diff --git a/datafusion/spark/src/function/string/mod.rs b/datafusion/spark/src/function/string/mod.rs index 9c90ded5f7e1b..bc94c27732c91 100644 --- a/datafusion/spark/src/function/string/mod.rs +++ b/datafusion/spark/src/function/string/mod.rs @@ -19,6 +19,7 @@ pub mod ascii; pub mod base64; pub mod char; pub mod concat; +pub mod concat_ws; pub mod elt; pub mod format_string; pub mod ilike; @@ -40,6 +41,7 @@ make_udf_function!(ascii::SparkAscii, ascii); make_udf_function!(base64::SparkBase64, base64); make_udf_function!(char::CharFunc, char); make_udf_function!(concat::SparkConcat, concat); +make_udf_function!(concat_ws::SparkConcatWs, concat_ws); make_udf_function!(ilike::SparkILike, ilike); make_udf_function!(length::SparkLengthFunc, length); make_udf_function!(elt::SparkElt, elt); @@ -77,6 +79,11 @@ pub mod expr_fn { "Concatenates multiple input strings into a single string. Returns NULL if any input is NULL.", args )); + export_functions!(( + concat_ws, + "Concatenates strings with separator. Supports arrays. Null values are skipped.", + sep args + )); export_functions!(( elt, "Returns the n-th input (1-indexed), e.g. returns 2nd input when n is 2. The function returns NULL if the index is 0 or exceeds the length of the array.", @@ -142,6 +149,7 @@ pub fn functions() -> Vec> { base64(), char(), concat(), + concat_ws(), elt(), ilike(), length(), diff --git a/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt b/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt index 62df636bba9ce..f6404cab8f3cd 100644 --- a/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt +++ b/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt @@ -21,22 +21,367 @@ # For more information, please see: # https://github.com/apache/datafusion/issues/15914 -## Original Query: SELECT concat_ws(' ', 'Spark', 'SQL'); -## PySpark 3.5.5 Result: {'concat_ws( , Spark, SQL)': 'Spark SQL', 'typeof(concat_ws( , Spark, SQL))': 'string', 'typeof( )': 'string', 'typeof(Spark)': 'string', 'typeof(SQL)': 'string'} -#query -#SELECT concat_ws(' '::string, 'Spark'::string, 'SQL'::string); - -## Original Query: SELECT concat_ws('/', 'foo', null, 'bar'); -## PySpark 3.5.5 Result: {'concat_ws(/, foo, NULL, bar)': 'foo/bar', 'typeof(concat_ws(/, foo, NULL, bar))': 'string', 'typeof(/)': 'string', 'typeof(foo)': 'string', 'typeof(NULL)': 'void', 'typeof(bar)': 'string'} -#query -#SELECT concat_ws('/'::string, 'foo'::string, NULL::void, 'bar'::string); - -## Original Query: SELECT concat_ws('s'); -## PySpark 3.5.5 Result: {'concat_ws(s)': '', 'typeof(concat_ws(s))': 'string', 'typeof(s)': 'string'} -#query -#SELECT concat_ws('s'::string); - -## Original Query: SELECT concat_ws(null, 'Spark', 'SQL'); -## PySpark 3.5.5 Result: {'concat_ws(NULL, Spark, SQL)': None, 'typeof(concat_ws(NULL, Spark, SQL))': 'string', 'typeof(NULL)': 'void', 'typeof(Spark)': 'string', 'typeof(SQL)': 'string'} -#query -#SELECT concat_ws(NULL::void, 'Spark'::string, 'SQL'::string); +## ── Basic scalar usage ────────────────────────────────────── + +## Multiple string arguments +query T +SELECT concat_ws(',', 'a', 'b', 'c'); +---- +a,b,c + +## Space separator +query T +SELECT concat_ws(' ', 'Spark', 'SQL'); +---- +Spark SQL + +## Slash separator with null skipped +query T +SELECT concat_ws('/', 'foo', NULL, 'bar'); +---- +foo/bar + +## Single argument after separator +query T +SELECT concat_ws(',', 'a'); +---- +a + +## No arguments after separator → empty string +query T +SELECT concat_ws(','); +---- +(empty) + +## Null separator returns null +query T +SELECT concat_ws(NULL, 'a', 'b', 'c'); +---- +NULL + +## All null arguments → empty string +query T +SELECT concat_ws(',', CAST(NULL AS STRING), CAST(NULL AS STRING)); +---- +(empty) + +## ── Array arguments ───────────────────────────────────────── + +## Array argument +query T +SELECT concat_ws(',', array('a', 'b', 'c')); +---- +a,b,c + +## Array with nulls skipped +query T +SELECT concat_ws(',', array('a', NULL, 'c')); +---- +a,c + +## Multiple arrays +query T +SELECT concat_ws(',', array('a', 'b'), array('c', 'd')); +---- +a,b,c,d + +## Mixed scalar and array arguments +query T +SELECT concat_ws(',', 'x', array('a', 'b'), 'y'); +---- +x,a,b,y + +## Null array is skipped +query T +SELECT concat_ws(',', 'x', CAST(NULL AS ARRAY), 'y'); +---- +x,y + +## ── Edge cases ─────────────────────────────────────────────── + +## Separator column with no value arguments +query T +SELECT concat_ws(sep) AS result FROM VALUES (','), ('-') AS t(sep); +---- +(empty) +(empty) + +## Null separator in column with no value arguments +query T +SELECT concat_ws(sep) AS result FROM VALUES (CAST(NULL AS STRING)), (',') AS t(sep); +---- +NULL +(empty) + +## ── Column expressions ────────────────────────────────────── + +## concat_ws on columns +query T +SELECT concat_ws('-', a, b) AS result FROM VALUES ('hello', 'world'), ('foo', 'bar') AS t(a, b); +---- +hello-world +foo-bar + +## concat_ws with null in columns +query T +SELECT concat_ws(',', a, b) AS result FROM VALUES ('a', 'b'), ('c', CAST(NULL AS STRING)), (CAST(NULL AS STRING), 'd') AS t(a, b); +---- +a,b +c +d + +## Scalar-only arguments over multiple rows (broadcast test) +query T +SELECT concat_ws(',', 'a', 'b') AS result FROM VALUES (1), (2), (3) AS t(x); +---- +a,b +a,b +a,b + +## ── Additional edge cases ─────────────────────────────────── + +## Empty separator — values concatenated with nothing between +query T +SELECT concat_ws('', 'a', 'b', 'c'); +---- +abc + +## Empty-string values are NOT skipped (only NULLs are) +query T +SELECT concat_ws(',', '', 'a', '', 'b'); +---- +,a,,b + +## Multi-character separator +query T +SELECT concat_ws(' - ', 'a', 'b', 'c'); +---- +a - b - c + +## Utf8View separator +query TT +SELECT concat_ws(arrow_cast(',', 'Utf8View'), 'a', 'b'), arrow_typeof(concat_ws(arrow_cast(',', 'Utf8View'), 'a', 'b')); +---- +a,b Utf8 + +## LargeUtf8 separator +query TT +SELECT concat_ws(arrow_cast(',', 'LargeUtf8'), 'a', 'b'), arrow_typeof(concat_ws(arrow_cast(',', 'LargeUtf8'), 'a', 'b')); +---- +a,b Utf8 + +## Empty array → empty string +query T +SELECT concat_ws(',', array()); +---- +(empty) + +## Scalar + array + array mix +query T +SELECT concat_ws(',', array('a', 'b'), 'c', array('d', 'e')); +---- +a,b,c,d,e + +## All-NULL row mixed with non-NULL rows +query T +SELECT concat_ws(',', a, b, c) AS result FROM VALUES + ('a', 'b', 'c'), + (CAST(NULL AS STRING), 'b', 'c'), + ('a', CAST(NULL AS STRING), CAST(NULL AS STRING)), + (CAST(NULL AS STRING), CAST(NULL AS STRING), CAST(NULL AS STRING)) + AS t(a, b, c); +---- +a,b,c +b,c +a +(empty) + +## Separator from column (per-row separator), with NULL rows +query T +SELECT concat_ws(sep, a, b) AS result FROM VALUES + (',', 'a', 'b'), + (CAST(NULL AS STRING), 'a', 'b'), + ('|', 'x', 'y') + AS t(sep, a, b); +---- +a,b +NULL +x|y + +## ── Spark cross-checked extras ────────────────────────────── + +## Zero arguments → error (Spark: WRONG_NUM_ARGS) +query error +SELECT concat_ws(); + +## Numeric separator coerced to string +query T +SELECT concat_ws(1, 'a', 'b'); +---- +a1b + +## Only numeric separator, no values → empty string +query T +SELECT concat_ws(123); +---- +(empty) + +## Numeric values coerced to string +query T +SELECT concat_ws(',', 1, 2, 3); +---- +1,2,3 + +## Float values +query T +SELECT concat_ws(',', 1.5, 2.5); +---- +1.5,2.5 + +## Boolean values +query T +SELECT concat_ws(',', true, false); +---- +true,false + +## Mixed numeric and string +query T +SELECT concat_ws(',', CAST(1 AS BIGINT), 'a'); +---- +1,a + +## Date values +query T +SELECT concat_ws(',', DATE '2024-01-01', 'x'); +---- +2024-01-01,x + +## Multi-byte UTF-8 separator +query T +SELECT concat_ws('é', 'a', 'b'); +---- +aéb + +## Nested concat_ws +query T +SELECT concat_ws('|', concat_ws(',', 'a', 'b'), concat_ws(',', 'c', 'd')); +---- +a,b|c,d + +## All-NULL elements in array → empty string +query T +SELECT concat_ws(',', array(CAST(NULL AS STRING), CAST(NULL AS STRING), CAST(NULL AS STRING))); +---- +(empty) + +## Empty-string elements in arrays are NOT skipped +query T +SELECT concat_ws(',', array(''), array('')); +---- +, + +## Multiple arrays interleaved with scalars and NULLs +query T +SELECT concat_ws(',', array('a', 'b'), 'c', CAST(NULL AS STRING), array('d'), '', 'e'); +---- +a,b,c,d,,e + +## Long argument list (variadic) +query T +SELECT concat_ws('-', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'); +---- +a-b-c-d-e-f-g-h-i-j + +## Long array +query T +SELECT concat_ws(',', array('a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j')); +---- +a,b,c,d,e,f,g,h,i,j + +## All-empty arguments +query T +SELECT concat_ws('', '', '', ''); +---- +(empty) + +## Empty separator with all NULLs +query T +SELECT concat_ws('', CAST(NULL AS STRING), CAST(NULL AS STRING)); +---- +(empty) + +## Long string preserved (length sanity) +query I +SELECT length(concat_ws(',', repeat('x', 1000), repeat('y', 1000))); +---- +2001 + +## ── List variants (FixedSizeList / ListView) ──────────────── + +## FixedSizeList argument is expanded element-by-element +query T +SELECT concat_ws(',', arrow_cast(make_array('a', 'b', 'c'), 'FixedSizeList(3, Utf8)')); +---- +a,b,c + +## ListView argument is expanded element-by-element +query T +SELECT concat_ws(',', arrow_cast(make_array('a', 'b'), 'ListView(Utf8)')); +---- +a,b + +## LargeListView argument is expanded element-by-element +query T +SELECT concat_ws(',', arrow_cast(make_array('a', 'b'), 'LargeListView(Utf8)')); +---- +a,b + +## ── Binary coercion (Spark casts binary to its UTF-8 view) ── + +## Binary argument is coerced to its string representation +query T +SELECT concat_ws(',', X'4869'); +---- +Hi + +## ── Null-separator-over-rows shape ────────────────────────── + +## Null separator on a multi-row column yields NULL on every row +query T +SELECT concat_ws(NULL, v) AS result FROM VALUES ('a'), ('b'), ('c') AS t(v) ORDER BY v; +---- +NULL +NULL +NULL + +## ── Non-string list elements (planner-inserted element cast) ─ + +## Array of integers — elements must be cast to STRING +query T +SELECT concat_ws(',', array(1, 2, 3)); +---- +1,2,3 + +## Array of doubles +query T +SELECT concat_ws('-', array(1.5, 2.5, 3.5)); +---- +1.5-2.5-3.5 + +## Array of booleans +query T +SELECT concat_ws(',', array(true, false, true)); +---- +true,false,true + +## Mixed: string scalar + int array + string scalar +query T +SELECT concat_ws(',', 'x', array(1, 2), 'y'); +---- +x,1,2,y + +## ── Struct rejection ──────────────────────────────────────── + +## Struct argument is rejected (not coerced to string) +query error +SELECT concat_ws(',', named_struct('a', 1)); From 574a1e6b39acd5b2c521d554863ac50fc0855654 Mon Sep 17 00:00:00 2001 From: "Ahmed EL." Date: Sat, 13 Jun 2026 02:40:47 +0100 Subject: [PATCH 4/4] fix: preserve Spark next_day whitespace validation (#22720) ## Which issue does this PR close? - Closes #22717. ## Rationale for this change Spark does not trim `dayOfWeek` before matching it in `next_day`, but `datafusion-spark` currently does. That makes values like `' MO '` succeed in DataFusion even though Spark treats them as invalid. ## What changes are included in this PR? - remove the `.trim()` call from `spark_next_day` - add a regression test proving whitespace-padded day names are rejected ## Are these changes tested? - `cargo test -p datafusion-spark next_day_rejects_whitespace_padded_day_names -- --nocapture` - `cargo test -p datafusion-spark` - `cargo fmt --all --check` - `cargo clippy -p datafusion-spark --all-targets --all-features --no-deps -- -D warnings` Note: the broader package clippy invocation still reports an existing unused import warning in untouched `datafusion/core/src/execution/session_state.rs` on current main. ## Are there any user-facing changes? Behavior now matches Spark for whitespace-padded `dayOfWeek` inputs in `next_day`. --------- Signed-off-by: xfocus3 Co-authored-by: xfocus3 Co-authored-by: Ahmed El amraouiyine --- datafusion/spark/src/function/datetime/next_day.rs | 8 +++++++- .../sqllogictest/test_files/spark/datetime/next_day.slt | 6 ++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/datafusion/spark/src/function/datetime/next_day.rs b/datafusion/spark/src/function/datetime/next_day.rs index 2241043d44cd7..2ef222526f387 100644 --- a/datafusion/spark/src/function/datetime/next_day.rs +++ b/datafusion/spark/src/function/datetime/next_day.rs @@ -210,7 +210,7 @@ where fn spark_next_day(days: i32, day_of_week: &str) -> Option { let date = Date32Type::to_naive_date_opt(days)?; - let day_of_week = day_of_week.trim().to_uppercase(); + 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"), @@ -279,4 +279,10 @@ mod tests { assert_eq!(field.data_type(), &DataType::Date32); assert!(field.is_nullable()); } + + #[test] + fn next_day_rejects_whitespace_padded_day_names() { + let monday = 19723; // 2024-01-01 + assert_eq!(spark_next_day(monday, " MO "), None); + } } diff --git a/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt b/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt index 872d1f2b58eb6..b0ffd7d0e412f 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt @@ -36,6 +36,12 @@ SELECT next_day('2015-07-27'::DATE, 'Sat'::string); ---- 2015-08-01 +# Whitespace-padded day names should be rejected (return NULL) per Spark behavior +query D +SELECT next_day('2015-01-14'::DATE, ' MO '::string); +---- +NULL + query error Failed to coerce arguments to satisfy a call to 'next_day' function SELECT next_day('2015-07-27'::DATE);