You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by co...@apache.org on 2025/09/02 19:47:46 UTC
(datafusion) branch main updated: feat(spark): implement Spark `width_bucket` function (#17331)
This is an automated email from the ASF dual-hosted git repository.
comphead pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 3f70c9a2b7 feat(spark): implement Spark `width_bucket` function (#17331)
3f70c9a2b7 is described below
commit 3f70c9a2b78a62bbb8eb4d96487318e709741b0a
Author: David López <da...@gmail.com>
AuthorDate: Tue Sep 2 21:47:40 2025 +0200
feat(spark): implement Spark `width_bucket` function (#17331)
* feat(spark): implement Spark width_bucket function
---
datafusion/spark/src/function/math/mod.rs | 13 +-
datafusion/spark/src/function/math/width_bucket.rs | 792 +++++++++++++++++++++
.../test_files/spark/math/width_bucket.slt | 126 +++-
3 files changed, 891 insertions(+), 40 deletions(-)
diff --git a/datafusion/spark/src/function/math/mod.rs b/datafusion/spark/src/function/math/mod.rs
index acc3479a08..092335e4aa 100644
--- a/datafusion/spark/src/function/math/mod.rs
+++ b/datafusion/spark/src/function/math/mod.rs
@@ -20,6 +20,7 @@ pub mod factorial;
pub mod hex;
pub mod modulus;
pub mod rint;
+pub mod width_bucket;
use datafusion_expr::ScalarUDF;
use datafusion_functions::make_udf_function;
@@ -31,6 +32,7 @@ make_udf_function!(hex::SparkHex, hex);
make_udf_function!(modulus::SparkMod, modulus);
make_udf_function!(modulus::SparkPmod, pmod);
make_udf_function!(rint::SparkRint, rint);
+make_udf_function!(width_bucket::SparkWidthBucket, width_bucket);
pub mod expr_fn {
use datafusion_functions::export_functions;
@@ -45,8 +47,17 @@ pub mod expr_fn {
export_functions!((modulus, "Returns the remainder of division of the first argument by the second argument.", arg1 arg2));
export_functions!((pmod, "Returns the positive remainder of division of the first argument by the second argument.", arg1 arg2));
export_functions!((rint, "Returns the double value that is closest in value to the argument and is equal to a mathematical integer.", arg1));
+ export_functions!((width_bucket, "Returns the bucket number into which the value of this expression would fall after being evaluated.", arg1 arg2 arg3 arg4));
}
pub fn functions() -> Vec<Arc<ScalarUDF>> {
- vec![expm1(), factorial(), hex(), modulus(), pmod(), rint()]
+ vec![
+ expm1(),
+ factorial(),
+ hex(),
+ modulus(),
+ pmod(),
+ rint(),
+ width_bucket(),
+ ]
}
diff --git a/datafusion/spark/src/function/math/width_bucket.rs b/datafusion/spark/src/function/math/width_bucket.rs
new file mode 100644
index 0000000000..24f8fe6b24
--- /dev/null
+++ b/datafusion/spark/src/function/math/width_bucket.rs
@@ -0,0 +1,792 @@
+// 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.
+
+use std::any::Any;
+use std::sync::Arc;
+
+use crate::function::error_utils::unsupported_data_types_exec_err;
+use arrow::array::{
+ Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalMonthDayNanoArray,
+ IntervalYearMonthArray,
+};
+use arrow::datatypes::DataType;
+use arrow::datatypes::DataType::{Duration, Float64, Int32, Interval};
+use arrow::datatypes::IntervalUnit::{MonthDayNano, YearMonth};
+use datafusion_common::cast::{
+ as_duration_microsecond_array, as_float64_array, as_int32_array,
+ as_interval_mdn_array, as_interval_ym_array,
+};
+use datafusion_common::{exec_err, Result};
+use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature};
+use datafusion_functions::utils::make_scalar_function;
+
+use arrow::array::{Int32Array, Int32Builder};
+use arrow::datatypes::TimeUnit::Microsecond;
+use datafusion_expr::Volatility::Immutable;
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkWidthBucket {
+ signature: Signature,
+}
+
+impl Default for SparkWidthBucket {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl SparkWidthBucket {
+ pub fn new() -> Self {
+ Self {
+ signature: Signature::user_defined(Immutable),
+ }
+ }
+}
+
+impl ScalarUDFImpl for SparkWidthBucket {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "width_bucket"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(Int32)
+ }
+
+ fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+ make_scalar_function(width_bucket_kern, vec![])(&args.args)
+ }
+
+ fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
+ if input.len() == 1 {
+ let value = &input[0];
+ Ok(value.sort_properties)
+ } else {
+ Ok(SortProperties::default())
+ }
+ }
+
+ fn coerce_types(&self, types: &[DataType]) -> Result<Vec<DataType>> {
+ use DataType::*;
+
+ let (v, lo, hi, n) = (&types[0], &types[1], &types[2], &types[3]);
+
+ let is_num = |t: &DataType| {
+ matches!(
+ t,
+ Int8 | Int16 | Int32 | Int64 | Float32 | Float64 | Decimal128(_, _)
+ )
+ };
+
+ match (v, lo, hi, n) {
+ (a, b, c, &(Int8 | Int16 | Int32 | Int64))
+ if is_num(a) && is_num(b) && is_num(c) =>
+ {
+ Ok(vec![Float64, Float64, Float64, Int32])
+ }
+ (
+ &Duration(_),
+ &Duration(_),
+ &Duration(_),
+ &(Int8 | Int16 | Int32 | Int64),
+ ) => Ok(vec![
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Duration(Microsecond),
+ Int32,
+ ]),
+ (
+ &Interval(MonthDayNano),
+ &Interval(MonthDayNano),
+ &Interval(MonthDayNano),
+ &(Int8 | Int16 | Int32 | Int64),
+ ) => Ok(vec![
+ Interval(MonthDayNano),
+ Interval(MonthDayNano),
+ Interval(MonthDayNano),
+ Int32,
+ ]),
+ (
+ &Interval(YearMonth),
+ &Interval(YearMonth),
+ &Interval(YearMonth),
+ &(Int8 | Int16 | Int32 | Int64),
+ ) => Ok(vec![
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Interval(YearMonth),
+ Int32,
+ ]),
+
+ _ => exec_err!(
+ "width_bucket expects a numeric argument, got {} {} {} {}",
+ types[0],
+ types[1],
+ types[2],
+ types[3]
+ ),
+ }
+ }
+}
+
+fn width_bucket_kern(args: &[ArrayRef]) -> Result<ArrayRef> {
+ let [v, minv, maxv, nb] = args else {
+ return exec_err!(
+ "width_bucket expects exactly 4 argument, got {}",
+ args.len()
+ );
+ };
+
+ match v.data_type() {
+ Float64 => {
+ let v = as_float64_array(v)?;
+ let min = as_float64_array(minv)?;
+ let max = as_float64_array(maxv)?;
+ let n_bucket = as_int32_array(nb)?;
+ Ok(Arc::new(width_bucket_float64(v, min, max, n_bucket)))
+ }
+ Duration(Microsecond) => {
+ let v = as_duration_microsecond_array(v)?;
+ let min = as_duration_microsecond_array(minv)?;
+ let max = as_duration_microsecond_array(maxv)?;
+ let n_bucket = as_int32_array(nb)?;
+ Ok(Arc::new(width_bucket_i64_as_float(v, min, max, n_bucket)))
+ }
+ Interval(YearMonth) => {
+ let v = as_interval_ym_array(v)?;
+ let min = as_interval_ym_array(minv)?;
+ let max = as_interval_ym_array(maxv)?;
+ let n_bucket = as_int32_array(nb)?;
+ Ok(Arc::new(width_bucket_i32_as_float(v, min, max, n_bucket)))
+ }
+ Interval(MonthDayNano) => {
+ let v = as_interval_mdn_array(v)?;
+ let min = as_interval_mdn_array(minv)?;
+ let max = as_interval_mdn_array(maxv)?;
+ let n_bucket = as_int32_array(nb)?;
+ Ok(Arc::new(width_bucket_interval_mdn_exact(v, min, max, n_bucket)))
+ }
+
+
+ other => Err(unsupported_data_types_exec_err(
+ "width_bucket",
+ "Float/Decimal OR Duration OR Interval(YearMonth) for first 3 args; Int for 4th",
+ &[
+ other.clone(),
+ minv.data_type().clone(),
+ maxv.data_type().clone(),
+ nb.data_type().clone(),
+ ],
+ )),
+ }
+}
+
+macro_rules! width_bucket_kernel_impl {
+ ($name:ident, $arr_ty:ty, $to_f64:expr, $check_nan:expr) => {
+ pub(crate) fn $name(
+ v: &$arr_ty,
+ min: &$arr_ty,
+ max: &$arr_ty,
+ n_bucket: &Int32Array,
+ ) -> Int32Array {
+ let len = v.len();
+ let mut b = Int32Builder::with_capacity(len);
+
+ for i in 0..len {
+ if v.is_null(i) || min.is_null(i) || max.is_null(i) || n_bucket.is_null(i)
+ {
+ b.append_null();
+ continue;
+ }
+ let x = ($to_f64)(v, i);
+ let l = ($to_f64)(min, i);
+ let h = ($to_f64)(max, i);
+ let buckets = n_bucket.value(i);
+
+ if buckets <= 0 {
+ b.append_null();
+ continue;
+ }
+ if $check_nan {
+ if !x.is_finite() || !l.is_finite() || !h.is_finite() {
+ b.append_null();
+ continue;
+ }
+ }
+
+ let ord = match l.partial_cmp(&h) {
+ Some(o) => o,
+ None => {
+ b.append_null();
+ continue;
+ }
+ };
+ if matches!(ord, std::cmp::Ordering::Equal) {
+ b.append_null();
+ continue;
+ }
+ let asc = matches!(ord, std::cmp::Ordering::Less);
+
+ if asc {
+ if x < l {
+ b.append_value(0);
+ continue;
+ }
+ if x >= h {
+ b.append_value(buckets + 1);
+ continue;
+ }
+ } else {
+ if x > l {
+ b.append_value(0);
+ continue;
+ }
+ if x <= h {
+ b.append_value(buckets + 1);
+ continue;
+ }
+ }
+
+ let width = (h - l) / (buckets as f64);
+ if width == 0.0 || !width.is_finite() {
+ b.append_null();
+ continue;
+ }
+ let mut bucket = ((x - l) / width).floor() as i32 + 1;
+ if bucket < 1 {
+ bucket = 1;
+ }
+ if bucket > buckets + 1 {
+ bucket = buckets + 1;
+ }
+
+ b.append_value(bucket);
+ }
+
+ b.finish()
+ }
+ };
+}
+
+width_bucket_kernel_impl!(
+ width_bucket_float64,
+ Float64Array,
+ |arr: &Float64Array, i: usize| arr.value(i),
+ true
+);
+
+width_bucket_kernel_impl!(
+ width_bucket_i64_as_float,
+ DurationMicrosecondArray,
+ |arr: &DurationMicrosecondArray, i: usize| arr.value(i) as f64,
+ false
+);
+
+width_bucket_kernel_impl!(
+ width_bucket_i32_as_float,
+ IntervalYearMonthArray,
+ |arr: &IntervalYearMonthArray, i: usize| arr.value(i) as f64,
+ false
+);
+const NS_PER_DAY_I128: i128 = 86_400_000_000_000;
+pub(crate) fn width_bucket_interval_mdn_exact(
+ v: &IntervalMonthDayNanoArray,
+ lo: &IntervalMonthDayNanoArray,
+ hi: &IntervalMonthDayNanoArray,
+ n: &Int32Array,
+) -> Int32Array {
+ let len = v.len();
+ let mut b = Int32Builder::with_capacity(len);
+
+ for i in 0..len {
+ if v.is_null(i) || lo.is_null(i) || hi.is_null(i) || n.is_null(i) {
+ b.append_null();
+ continue;
+ }
+ let buckets = n.value(i);
+ if buckets <= 0 {
+ b.append_null();
+ continue;
+ }
+
+ let x = v.value(i);
+ let l = lo.value(i);
+ let h = hi.value(i);
+
+ // asc/desc
+ // Values of IntervalMonthDayNano are compared using their binary representation, which can lead to surprising results.
+ let asc = (l.months, l.days, l.nanoseconds) < (h.months, h.days, h.nanoseconds);
+ if (l.months, l.days, l.nanoseconds) == (h.months, h.days, h.nanoseconds) {
+ b.append_null();
+ continue;
+ }
+
+ // ------------------- only month -------------------
+ if l.days == h.days && l.nanoseconds == h.nanoseconds && l.months != h.months {
+ let x_m = x.months as f64;
+ let l_m = l.months as f64;
+ let h_m = h.months as f64;
+
+ if asc {
+ if x_m < l_m {
+ b.append_value(0);
+ continue;
+ }
+ if x_m >= h_m {
+ b.append_value(buckets + 1);
+ continue;
+ }
+ } else {
+ if x_m > l_m {
+ b.append_value(0);
+ continue;
+ }
+ if x_m <= h_m {
+ b.append_value(buckets + 1);
+ continue;
+ }
+ }
+
+ let width = (h_m - l_m) / (buckets as f64);
+ if width == 0.0 || !width.is_finite() {
+ b.append_null();
+ continue;
+ }
+
+ let mut bucket = ((x_m - l_m) / width).floor() as i32 + 1;
+ if bucket < 1 {
+ bucket = 1;
+ }
+ if bucket > buckets + 1 {
+ bucket = buckets + 1;
+ }
+ b.append_value(bucket);
+ continue;
+ }
+
+ // --------------- months equals -------------------
+ if l.months == h.months {
+ let base_days = l.days as i128;
+ let base_ns = l.nanoseconds as i128;
+
+ let xf = (x.days as i128 - base_days) * NS_PER_DAY_I128
+ + (x.nanoseconds as i128 - base_ns);
+ let hf = (h.days as i128 - base_days) * NS_PER_DAY_I128
+ + (h.nanoseconds as i128 - base_ns);
+
+ let x_f = xf as f64;
+ let l_f = 0.0;
+ let h_f = hf as f64;
+
+ if asc {
+ if x_f < l_f {
+ b.append_value(0);
+ continue;
+ }
+ if x_f >= h_f {
+ b.append_value(buckets + 1);
+ continue;
+ }
+ } else {
+ if x_f > l_f {
+ b.append_value(0);
+ continue;
+ }
+ if x_f <= h_f {
+ b.append_value(buckets + 1);
+ continue;
+ }
+ }
+
+ let width = (h_f - l_f) / (buckets as f64);
+ if width == 0.0 || !width.is_finite() {
+ b.append_null();
+ continue;
+ }
+
+ let mut bucket = ((x_f - l_f) / width).floor() as i32 + 1;
+ if bucket < 1 {
+ bucket = 1;
+ }
+ if bucket > buckets + 1 {
+ bucket = buckets + 1;
+ }
+ b.append_value(bucket);
+ continue;
+ }
+
+ b.append_null();
+ }
+
+ b.finish()
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::sync::Arc;
+
+ use arrow::array::{
+ ArrayRef, DurationMicrosecondArray, Float64Array, Int32Array,
+ IntervalYearMonthArray,
+ };
+ use arrow::datatypes::IntervalMonthDayNano;
+
+ // --- Helpers -------------------------------------------------------------
+
+ fn i32_array_all(len: usize, val: i32) -> Arc<Int32Array> {
+ Arc::new(Int32Array::from(vec![val; len]))
+ }
+
+ fn f64_array(vals: &[f64]) -> Arc<Float64Array> {
+ Arc::new(Float64Array::from(vals.to_vec()))
+ }
+
+ fn f64_array_opt(vals: &[Option<f64>]) -> Arc<Float64Array> {
+ Arc::new(Float64Array::from(vals.to_vec()))
+ }
+
+ fn dur_us_array(vals: &[i64]) -> Arc<DurationMicrosecondArray> {
+ Arc::new(DurationMicrosecondArray::from(vals.to_vec()))
+ }
+
+ fn ym_array(vals: &[i32]) -> Arc<IntervalYearMonthArray> {
+ Arc::new(IntervalYearMonthArray::from(vals.to_vec()))
+ }
+
+ fn downcast_i32(arr: &ArrayRef) -> &Int32Array {
+ arr.as_any().downcast_ref::<Int32Array>().unwrap()
+ }
+
+ fn mdn_array(vals: &[(i32, i32, i64)]) -> Arc<IntervalMonthDayNanoArray> {
+ let data: Vec<IntervalMonthDayNano> = vals
+ .iter()
+ .map(|(m, d, ns)| IntervalMonthDayNano::new(*m, *d, *ns))
+ .collect();
+ Arc::new(IntervalMonthDayNanoArray::from(data))
+ }
+
+ // --- Float64 -------------------------------------------------------------
+
+ #[test]
+ fn test_width_bucket_f64_basic() {
+ let v = f64_array(&[0.5, 1.0, 9.9, -1.0, 10.0]);
+ let lo = f64_array(&[0.0, 0.0, 0.0, 0.0, 0.0]);
+ let hi = f64_array(&[10.0, 10.0, 10.0, 10.0, 10.0]);
+ let n = i32_array_all(5, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert_eq!(out.values(), &[1, 2, 10, 0, 11]);
+ }
+
+ #[test]
+ fn test_width_bucket_f64_descending_range() {
+ let v = f64_array(&[9.9, 10.0, 0.0, -0.1, 10.1]);
+ let lo = f64_array(&[10.0; 5]);
+ let hi = f64_array(&[0.0; 5]);
+ let n = i32_array_all(5, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+
+ assert_eq!(out.values(), &[1, 1, 11, 11, 0]);
+ }
+ #[test]
+ fn test_width_bucket_f64_bounds_inclusive_exclusive_asc() {
+ let v = f64_array(&[0.0, 9.999999999, 10.0]);
+ let lo = f64_array(&[0.0; 3]);
+ let hi = f64_array(&[10.0; 3]);
+ let n = i32_array_all(3, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert_eq!(out.values(), &[1, 10, 11]);
+ }
+
+ #[test]
+ fn test_width_bucket_f64_bounds_inclusive_exclusive_desc() {
+ let v = f64_array(&[10.0, 0.0, -0.000001]);
+ let lo = f64_array(&[10.0; 3]);
+ let hi = f64_array(&[0.0; 3]);
+ let n = i32_array_all(3, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert_eq!(out.values(), &[1, 11, 11]);
+ }
+
+ #[test]
+ fn test_width_bucket_f64_edge_cases() {
+ let v = f64_array(&[1.0, 5.0, 9.0]);
+ let lo = f64_array(&[0.0, 0.0, 0.0]);
+ let hi = f64_array(&[10.0, 10.0, 10.0]);
+ let n = Arc::new(Int32Array::from(vec![0, -1, 10]));
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert!(out.is_null(0));
+ assert!(out.is_null(1));
+ assert_eq!(out.value(2), 10);
+
+ let v = f64_array(&[1.0]);
+ let lo = f64_array(&[5.0]);
+ let hi = f64_array(&[5.0]);
+ let n = i32_array_all(1, 10);
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert!(out.is_null(0));
+
+ let v = f64_array_opt(&[Some(f64::NAN)]);
+ let lo = f64_array(&[0.0]);
+ let hi = f64_array(&[10.0]);
+ let n = i32_array_all(1, 10);
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert!(out.is_null(0));
+ }
+
+ #[test]
+ fn test_width_bucket_f64_nulls_propagate() {
+ let v = f64_array_opt(&[None, Some(1.0), Some(2.0), Some(3.0)]);
+ let lo = f64_array(&[0.0; 4]);
+ let hi = f64_array(&[10.0; 4]);
+ let n = i32_array_all(4, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert!(out.is_null(0));
+ assert_eq!(out.value(1), 2);
+ assert_eq!(out.value(2), 3);
+ assert_eq!(out.value(3), 4);
+
+ let v = f64_array(&[1.0]);
+ let lo = f64_array_opt(&[None]);
+ let hi = f64_array(&[10.0]);
+ let n = i32_array_all(1, 10);
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert!(out.is_null(0));
+ }
+
+ // --- Duration(Microsecond) ----------------------------------------------
+
+ #[test]
+ fn test_width_bucket_duration_us() {
+ let v = dur_us_array(&[1_000_000, 0, -1]);
+ let lo = dur_us_array(&[0, 0, 0]);
+ let hi = dur_us_array(&[2_000_000, 2_000_000, 2_000_000]);
+ let n = i32_array_all(3, 2);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert_eq!(out.values(), &[2, 1, 0]);
+ }
+
+ #[test]
+ fn test_width_bucket_duration_us_equal_bounds() {
+ let v = dur_us_array(&[0]);
+ let lo = dur_us_array(&[1]);
+ let hi = dur_us_array(&[1]);
+ let n = i32_array_all(1, 10);
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ assert!(downcast_i32(&out).is_null(0));
+ }
+
+ // --- Interval(YearMonth) ------------------------------------------------
+
+ #[test]
+ fn test_width_bucket_interval_ym_basic() {
+ let v = ym_array(&[0, 5, 11, 12, 13]);
+ let lo = ym_array(&[0; 5]);
+ let hi = ym_array(&[12; 5]);
+ let n = i32_array_all(5, 12);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert_eq!(out.values(), &[1, 6, 12, 13, 13]);
+ }
+
+ #[test]
+ fn test_width_bucket_interval_ym_desc() {
+ let v = ym_array(&[11, 12, 0, -1, 13]);
+ let lo = ym_array(&[12; 5]);
+ let hi = ym_array(&[0; 5]);
+ let n = i32_array_all(5, 12);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert_eq!(out.values(), &[2, 1, 13, 13, 0]);
+ }
+
+ // --- Interval(MonthDayNano) --------------------------------------------
+
+ #[test]
+ fn test_width_bucket_interval_mdn_months_only_basic() {
+ let v = mdn_array(&[(0, 0, 0), (5, 0, 0), (11, 0, 0), (12, 0, 0), (13, 0, 0)]);
+ let lo = mdn_array(&[(0, 0, 0); 5]);
+ let hi = mdn_array(&[(12, 0, 0); 5]);
+ let n = i32_array_all(5, 12);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert_eq!(out.values(), &[1, 6, 12, 13, 13]);
+ }
+
+ #[test]
+ fn test_width_bucket_interval_mdn_months_only_desc() {
+ let v = mdn_array(&[(11, 0, 0), (12, 0, 0), (0, 0, 0), (-1, 0, 0), (13, 0, 0)]);
+ let lo = mdn_array(&[(12, 0, 0); 5]);
+ let hi = mdn_array(&[(0, 0, 0); 5]);
+ let n = i32_array_all(5, 12);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ // Mismo patrón que YM descendente
+ assert_eq!(out.values(), &[2, 1, 13, 13, 0]);
+ }
+
+ #[test]
+ fn test_width_bucket_interval_mdn_day_nano_basic() {
+ let v = mdn_array(&[
+ (0, 0, 0),
+ (0, 5, 0),
+ (0, 9, 0),
+ (0, 10, 0),
+ (0, -1, 0),
+ (0, 11, 0),
+ ]);
+ let lo = mdn_array(&[(0, 0, 0); 6]);
+ let hi = mdn_array(&[(0, 10, 0); 6]);
+ let n = i32_array_all(6, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ // x==hi -> n+1, x<lo -> 0, x>hi -> n+1
+ assert_eq!(out.values(), &[1, 6, 10, 11, 0, 11]);
+ }
+
+ #[test]
+ fn test_width_bucket_interval_mdn_day_nano_desc() {
+ let v = mdn_array(&[(0, 9, 0), (0, 10, 0), (0, 0, 0), (0, -1, 0), (0, 11, 0)]);
+ let lo = mdn_array(&[(0, 10, 0); 5]);
+ let hi = mdn_array(&[(0, 0, 0); 5]);
+ let n = i32_array_all(5, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+
+ assert_eq!(out.values(), &[2, 1, 11, 11, 0]);
+ }
+ #[test]
+ fn test_width_bucket_interval_mdn_day_nano_desc_inside() {
+ let v = mdn_array(&[(0, 9, 1), (0, 10, 0), (0, 0, 0), (0, -1, 0), (0, 11, 0)]);
+ let lo = mdn_array(&[(0, 10, 0); 5]);
+ let hi = mdn_array(&[(0, 0, 0); 5]);
+ let n = i32_array_all(5, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+
+ assert_eq!(out.values(), &[1, 1, 11, 11, 0]);
+ }
+
+ #[test]
+ fn test_width_bucket_interval_mdn_mixed_months_and_days_is_null() {
+ let v = mdn_array(&[(0, 1, 0)]);
+ let lo = mdn_array(&[(0, 0, 0)]);
+ let hi = mdn_array(&[(1, 1, 0)]);
+ let n = i32_array_all(1, 4);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert!(out.is_null(0));
+ }
+
+ #[test]
+ fn test_width_bucket_interval_mdn_equal_bounds_is_null() {
+ let v = mdn_array(&[(0, 0, 0)]);
+ let lo = mdn_array(&[(1, 2, 3)]);
+ let hi = mdn_array(&[(1, 2, 3)]); // lo == hi
+ let n = i32_array_all(1, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ assert!(downcast_i32(&out).is_null(0));
+ }
+
+ #[test]
+ fn test_width_bucket_interval_mdn_invalid_n_is_null() {
+ let v = mdn_array(&[(0, 0, 0)]);
+ let lo = mdn_array(&[(0, 0, 0)]);
+ let hi = mdn_array(&[(0, 10, 0)]);
+ let n = Arc::new(Int32Array::from(vec![0])); // n <= 0
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ assert!(downcast_i32(&out).is_null(0));
+ }
+
+ #[test]
+ fn test_width_bucket_interval_mdn_nulls_propagate() {
+ let v = Arc::new(IntervalMonthDayNanoArray::from(vec![
+ None,
+ Some(IntervalMonthDayNano::new(0, 5, 0)),
+ ]));
+ let lo = mdn_array(&[(0, 0, 0), (0, 0, 0)]);
+ let hi = mdn_array(&[(0, 10, 0), (0, 10, 0)]);
+ let n = i32_array_all(2, 10);
+
+ let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
+ let out = downcast_i32(&out);
+ assert!(out.is_null(0));
+ assert_eq!(out.value(1), 6);
+ }
+
+ // --- Errores -------------------------------------------------------------
+
+ #[test]
+ fn test_width_bucket_wrong_arg_count() {
+ let v = f64_array(&[1.0]);
+ let lo = f64_array(&[0.0]);
+ let hi = f64_array(&[10.0]);
+ let err = width_bucket_kern(&[v, lo, hi]).unwrap_err();
+ let msg = format!("{err}");
+ assert!(msg.contains("expects exactly 4"), "unexpected error: {msg}");
+ }
+
+ #[test]
+ fn test_width_bucket_unsupported_type() {
+ let v: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
+ let lo = f64_array(&[0.0, 0.0, 0.0]);
+ let hi = f64_array(&[10.0, 10.0, 10.0]);
+ let n = i32_array_all(3, 10);
+
+ let err = width_bucket_kern(&[v, lo, hi, n]).unwrap_err();
+ let msg = format!("{err}");
+ assert!(
+ msg.contains("unsupported data types")
+ || msg.contains("Float/Decimal OR Duration OR Interval(YearMonth)"),
+ "unexpected error: {msg}"
+ );
+ }
+}
diff --git a/datafusion/sqllogictest/test_files/spark/math/width_bucket.slt b/datafusion/sqllogictest/test_files/spark/math/width_bucket.slt
index a463f430a7..d2661ceb9d 100644
--- a/datafusion/sqllogictest/test_files/spark/math/width_bucket.slt
+++ b/datafusion/sqllogictest/test_files/spark/math/width_bucket.slt
@@ -21,42 +21,90 @@
# For more information, please see:
# https://github.com/apache/datafusion/issues/15914
-## Original Query: SELECT width_bucket(-0.9, 5.2, 0.5, 2);
-## PySpark 3.5.5 Result: {'width_bucket(-0.9, 5.2, 0.5, 2)': 3, 'typeof(width_bucket(-0.9, 5.2, 0.5, 2))': 'bigint', 'typeof(-0.9)': 'decimal(1,1)', 'typeof(5.2)': 'decimal(2,1)', 'typeof(0.5)': 'decimal(1,1)', 'typeof(2)': 'int'}
-#query
-#SELECT width_bucket(-0.9::decimal(1,1), 5.2::decimal(2,1), 0.5::decimal(1,1), 2::int);
-
-## Original Query: SELECT width_bucket(-2.1, 1.3, 3.4, 3);
-## PySpark 3.5.5 Result: {'width_bucket(-2.1, 1.3, 3.4, 3)': 0, 'typeof(width_bucket(-2.1, 1.3, 3.4, 3))': 'bigint', 'typeof(-2.1)': 'decimal(2,1)', 'typeof(1.3)': 'decimal(2,1)', 'typeof(3.4)': 'decimal(2,1)', 'typeof(3)': 'int'}
-#query
-#SELECT width_bucket(-2.1::decimal(2,1), 1.3::decimal(2,1), 3.4::decimal(2,1), 3::int);
-
-## Original Query: SELECT width_bucket(5.3, 0.2, 10.6, 5);
-## PySpark 3.5.5 Result: {'width_bucket(5.3, 0.2, 10.6, 5)': 3, 'typeof(width_bucket(5.3, 0.2, 10.6, 5))': 'bigint', 'typeof(5.3)': 'decimal(2,1)', 'typeof(0.2)': 'decimal(1,1)', 'typeof(10.6)': 'decimal(3,1)', 'typeof(5)': 'int'}
-#query
-#SELECT width_bucket(5.3::decimal(2,1), 0.2::decimal(1,1), 10.6::decimal(3,1), 5::int);
-
-## Original Query: SELECT width_bucket(8.1, 0.0, 5.7, 4);
-## PySpark 3.5.5 Result: {'width_bucket(8.1, 0.0, 5.7, 4)': 5, 'typeof(width_bucket(8.1, 0.0, 5.7, 4))': 'bigint', 'typeof(8.1)': 'decimal(2,1)', 'typeof(0.0)': 'decimal(1,1)', 'typeof(5.7)': 'decimal(2,1)', 'typeof(4)': 'int'}
-#query
-#SELECT width_bucket(8.1::decimal(2,1), 0.0::decimal(1,1), 5.7::decimal(2,1), 4::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 1, "typeof(width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 1, "typeof(width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)": 2, "typeof(width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10))": 'bigint', "typeof(INTERVAL '1' DAY)": 'interval day', "typeof(INTERVAL '0' DAY)": 'interval day', "typeof(INTERVAL '10' DAY)": 'interval day', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' DAY::interval day, INTERVAL '0' DAY::interval day, INTERVAL '10' DAY::interval day, 10::int);
-
-## Original Query: SELECT width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10);
-## PySpark 3.5.5 Result: {"width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)": 2, "typeof(width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10))": 'bigint', "typeof(INTERVAL '1' YEAR)": 'interval year', "typeof(INTERVAL '0' YEAR)": 'interval year', "typeof(INTERVAL '10' YEAR)": 'interval year', 'typeof(10)': 'int'}
-#query
-#SELECT width_bucket(INTERVAL '1' YEAR::interval year, INTERVAL '0' YEAR::interval year, INTERVAL '10' YEAR::interval year, 10::int);
+query I
+SELECT width_bucket(-0.9, 5.2, 0.5, 2)
+----
+3
+
+query I
+SELECT width_bucket(-2.1, 1.3, 3.4, 3)
+----
+0
+
+query I
+SELECT width_bucket(5.3, 0.2, 10.6, 5)
+----
+3
+
+query I
+SELECT width_bucket(8.1, 0.0, 5.7, 4)
+----
+5
+
+query I
+SELECT width_bucket(INTERVAL '0' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)
+----
+1
+
+query I
+SELECT width_bucket(INTERVAL '0' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)
+----
+1
+
+query I
+SELECT width_bucket(INTERVAL '1' DAY, INTERVAL '0' DAY, INTERVAL '10' DAY, 10)
+----
+2
+
+query I
+SELECT width_bucket(INTERVAL '1' YEAR, INTERVAL '0' YEAR, INTERVAL '10' YEAR, 10)
+----
+2
+
+# test of sail
+query I
+SELECT width_bucket(0.0, 10.0, 0.0, 5)
+----
+6
+
+query I
+SELECT width_bucket(10.0, 0.0, 10.0, 5)
+----
+6
+
+query I
+SELECT width_bucket(10.0, 0.0, 0.0, 5)
+----
+NULL
+
+# lo == hi
+query I
+SELECT width_bucket(10.0, 0.0, 0.0, 5);
+----
+NULL
+
+# n <= 0
+query I
+SELECT width_bucket(5.0, 0.0, 10.0, 0);
+----
+NULL
+
+query I
+SELECT width_bucket(arrow_cast('NaN','Float64'),5.0, 0.0, 5)
+----
+NULL
+
+query I
+SELECT width_bucket(5.0, arrow_cast('NaN','Float64'), 0.0, 5)
+----
+NULL
+
+query I
+SELECT width_bucket(5.0, 0.0, arrow_cast('NaN','Float64'), 5)
+----
+NULL
+
+query I
+SELECT width_bucket(INTERVAL '1' YEAR, INTERVAL '5' YEAR, INTERVAL '5' YEAR, 10)
+----
+NULL
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org