You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by mb...@apache.org on 2026/06/22 15:28:36 UTC
(datafusion) branch branch-54 updated: [branch-54]: backport #22811 (bugfix: changed return type of spark's width_bucket to i64) (#23087)
This is an automated email from the ASF dual-hosted git repository.
mbutrovich pushed a commit to branch branch-54
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/branch-54 by this push:
new aa59c66c10 [branch-54]: backport #22811 (bugfix: changed return type of spark's width_bucket to i64) (#23087)
aa59c66c10 is described below
commit aa59c66c10e58d48e60ee30f19f27a02eb8a8577
Author: Matt Butrovich <mb...@users.noreply.github.com>
AuthorDate: Mon Jun 22 11:28:30 2026 -0400
[branch-54]: backport #22811 (bugfix: changed return type of spark's width_bucket to i64) (#23087)
backport #22811 to `branch-54` in anticipation of 54.1.0 release.
DataFusion Comet encountered this issue while trying to upgrade to DF
54.0.0, so it would ease that upgrade for Comet if this were backported.
Co-authored-by: Eduardo Aguilar <co...@eduaguilar.com>
---
datafusion/spark/src/function/math/width_bucket.rs | 78 ++++++++++------------
1 file changed, 37 insertions(+), 41 deletions(-)
diff --git a/datafusion/spark/src/function/math/width_bucket.rs b/datafusion/spark/src/function/math/width_bucket.rs
index 79da924116..93be47a4ca 100644
--- a/datafusion/spark/src/function/math/width_bucket.rs
+++ b/datafusion/spark/src/function/math/width_bucket.rs
@@ -21,8 +21,7 @@ use arrow::array::{
Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalMonthDayNanoArray,
IntervalYearMonthArray,
};
-use arrow::datatypes::DataType;
-use arrow::datatypes::DataType::{Duration, Float64, Int32, Interval};
+use arrow::datatypes::DataType::{self, Duration, Float64, Int64, Interval};
use arrow::datatypes::IntervalUnit::{MonthDayNano, YearMonth};
use datafusion_common::cast::{
as_duration_microsecond_array, as_float64_array, as_int64_array,
@@ -40,7 +39,7 @@ use datafusion_expr::{
};
use datafusion_functions::utils::make_scalar_function;
-use arrow::array::{Int32Array, Int32Builder, Int64Array};
+use arrow::array::{Int64Array, Int64Builder};
use arrow::datatypes::TimeUnit::Microsecond;
use datafusion_expr::Coercion;
use datafusion_expr::Volatility::Immutable;
@@ -125,7 +124,7 @@ impl ScalarUDFImpl for SparkWidthBucket {
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
- Ok(Int32)
+ Ok(Int64)
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
@@ -199,9 +198,9 @@ macro_rules! width_bucket_kernel_impl {
min: &$arr_ty,
max: &$arr_ty,
n_bucket: &Int64Array,
- ) -> Int32Array {
+ ) -> Int64Array {
let len = v.len();
- let mut b = Int32Builder::with_capacity(len);
+ let mut b = Int64Builder::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)
@@ -218,7 +217,7 @@ macro_rules! width_bucket_kernel_impl {
b.append_null();
continue;
}
- let next_bucket = (buckets + 1) as i32;
+ let next_bucket = (buckets + 1) as i64;
if $check_nan {
if !x.is_finite() || !l.is_finite() || !h.is_finite() {
b.append_null();
@@ -264,7 +263,7 @@ macro_rules! width_bucket_kernel_impl {
b.append_null();
continue;
}
- let mut bucket = ((x - l) / width).floor() as i32 + 1;
+ let mut bucket = ((x - l) / width).floor() as i64 + 1;
if bucket < 1 {
bucket = 1;
}
@@ -306,9 +305,9 @@ pub(crate) fn width_bucket_interval_mdn_exact(
lo: &IntervalMonthDayNanoArray,
hi: &IntervalMonthDayNanoArray,
n: &Int64Array,
-) -> Int32Array {
+) -> Int64Array {
let len = v.len();
- let mut b = Int32Builder::with_capacity(len);
+ let mut b = Int64Builder::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) {
@@ -320,7 +319,7 @@ pub(crate) fn width_bucket_interval_mdn_exact(
b.append_null();
continue;
}
- let next_bucket = (buckets + 1) as i32;
+ let next_bucket = buckets + 1;
let x = v.value(i);
let l = lo.value(i);
@@ -366,7 +365,7 @@ pub(crate) fn width_bucket_interval_mdn_exact(
continue;
}
- let mut bucket = ((x_m - l_m) / width).floor() as i32 + 1;
+ let mut bucket = ((x_m - l_m) / width).floor() as i64 + 1;
if bucket < 1 {
bucket = 1;
}
@@ -417,7 +416,7 @@ pub(crate) fn width_bucket_interval_mdn_exact(
continue;
}
- let mut bucket = ((x_f - l_f) / width).floor() as i32 + 1;
+ let mut bucket = ((x_f - l_f) / width).floor() as i64 + 1;
if bucket < 1 {
bucket = 1;
}
@@ -437,10 +436,11 @@ pub(crate) fn width_bucket_interval_mdn_exact(
#[cfg(test)]
mod tests {
use super::*;
+ use arrow::datatypes::Int64Type;
use arrow::array::{
- ArrayRef, DurationMicrosecondArray, Float64Array, Int32Array, Int64Array,
- IntervalYearMonthArray,
+ ArrayRef, AsArray, DurationMicrosecondArray, Float64Array, Int32Array,
+ Int64Array, IntervalYearMonthArray,
};
use arrow::datatypes::IntervalMonthDayNano;
@@ -466,10 +466,6 @@ mod tests {
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()
@@ -488,7 +484,7 @@ mod tests {
let n = i64_array_all(5, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[1, 2, 10, 0, 11]);
}
@@ -500,7 +496,7 @@ mod tests {
let n = i64_array_all(5, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[1, 1, 11, 11, 0]);
}
@@ -512,7 +508,7 @@ mod tests {
let n = i64_array_all(3, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[1, 10, 11]);
}
@@ -524,7 +520,7 @@ mod tests {
let n = i64_array_all(3, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[1, 11, 11]);
}
@@ -535,7 +531,7 @@ mod tests {
let hi = f64_array(&[10.0, 10.0, 10.0]);
let n = Arc::new(Int64Array::from(vec![0, -1, 10]));
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert!(out.is_null(0));
assert!(out.is_null(1));
assert_eq!(out.value(2), 10);
@@ -545,7 +541,7 @@ mod tests {
let hi = f64_array(&[5.0]);
let n = i64_array_all(1, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert!(out.is_null(0));
let v = f64_array_opt(&[Some(f64::NAN)]);
@@ -553,7 +549,7 @@ mod tests {
let hi = f64_array(&[10.0]);
let n = i64_array_all(1, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert!(out.is_null(0));
}
@@ -565,7 +561,7 @@ mod tests {
let n = i64_array_all(4, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert!(out.is_null(0));
assert_eq!(out.value(1), 2);
assert_eq!(out.value(2), 3);
@@ -576,7 +572,7 @@ mod tests {
let hi = f64_array(&[10.0]);
let n = i64_array_all(1, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert!(out.is_null(0));
}
@@ -590,7 +586,7 @@ mod tests {
let n = i64_array_all(3, 2);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[2, 1, 0]);
}
@@ -601,7 +597,7 @@ mod tests {
let hi = dur_us_array(&[1]);
let n = i64_array_all(1, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- assert!(downcast_i32(&out).is_null(0));
+ assert!(out.as_primitive::<Int64Type>().is_null(0));
}
// --- Interval(YearMonth) ------------------------------------------------
@@ -614,7 +610,7 @@ mod tests {
let n = i64_array_all(5, 12);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[1, 6, 12, 13, 13]);
}
@@ -626,7 +622,7 @@ mod tests {
let n = i64_array_all(5, 12);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[2, 1, 13, 13, 0]);
}
@@ -640,7 +636,7 @@ mod tests {
let n = i64_array_all(5, 12);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[1, 6, 12, 13, 13]);
}
@@ -652,7 +648,7 @@ mod tests {
let n = i64_array_all(5, 12);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
// Mismo patrón que YM descendente
assert_eq!(out.values(), &[2, 1, 13, 13, 0]);
}
@@ -672,7 +668,7 @@ mod tests {
let n = i64_array_all(6, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
// x==hi -> n+1, x<lo -> 0, x>hi -> n+1
assert_eq!(out.values(), &[1, 6, 10, 11, 0, 11]);
}
@@ -685,7 +681,7 @@ mod tests {
let n = i64_array_all(5, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[2, 1, 11, 11, 0]);
}
@@ -697,7 +693,7 @@ mod tests {
let n = i64_array_all(5, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert_eq!(out.values(), &[1, 1, 11, 11, 0]);
}
@@ -710,7 +706,7 @@ mod tests {
let n = i64_array_all(1, 4);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert!(out.is_null(0));
}
@@ -722,7 +718,7 @@ mod tests {
let n = i64_array_all(1, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- assert!(downcast_i32(&out).is_null(0));
+ assert!(out.as_primitive::<Int64Type>().is_null(0));
}
#[test]
@@ -733,7 +729,7 @@ mod tests {
let n = Arc::new(Int64Array::from(vec![0])); // n <= 0
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- assert!(downcast_i32(&out).is_null(0));
+ assert!(out.as_primitive::<Int64Type>().is_null(0));
}
#[test]
@@ -747,7 +743,7 @@ mod tests {
let n = i64_array_all(2, 10);
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
- let out = downcast_i32(&out);
+ let out = out.as_primitive::<Int64Type>();
assert!(out.is_null(0));
assert_eq!(out.value(1), 6);
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org