Skip to content

Commit aa59c66

Browse files
mbutrovichaguilaredu
mbutrovich
and
aguilaredu
authored
[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 <contact@eduaguilar.com>
1 parent 08da279 commit aa59c66

1 file changed

Lines changed: 37 additions & 41 deletions

File tree

‎datafusion/spark/src/function/math/width_bucket.rs‎

Lines changed: 37 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,7 @@ use arrow::array::{
2121
Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalMonthDayNanoArray,
2222
IntervalYearMonthArray,
2323
};
24-
use arrow::datatypes::DataType;
25-
use arrow::datatypes::DataType::{Duration, Float64, Int32, Interval};
24+
use arrow::datatypes::DataType::{self, Duration, Float64, Int64, Interval};
2625
use arrow::datatypes::IntervalUnit::{MonthDayNano, YearMonth};
2726
use datafusion_common::cast::{
2827
as_duration_microsecond_array, as_float64_array, as_int64_array,
@@ -40,7 +39,7 @@ use datafusion_expr::{
4039
};
4140
use datafusion_functions::utils::make_scalar_function;
4241

43-
use arrow::array::{Int32Array, Int32Builder, Int64Array};
42+
use arrow::array::{Int64Array, Int64Builder};
4443
use arrow::datatypes::TimeUnit::Microsecond;
4544
use datafusion_expr::Coercion;
4645
use datafusion_expr::Volatility::Immutable;
@@ -125,7 +124,7 @@ impl ScalarUDFImpl for SparkWidthBucket {
125124
}
126125

127126
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
128-
Ok(Int32)
127+
Ok(Int64)
129128
}
130129

131130
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
@@ -199,9 +198,9 @@ macro_rules! width_bucket_kernel_impl {
199198
min: &$arr_ty,
200199
max: &$arr_ty,
201200
n_bucket: &Int64Array,
202-
) -> Int32Array {
201+
) -> Int64Array {
203202
let len = v.len();
204-
let mut b = Int32Builder::with_capacity(len);
203+
let mut b = Int64Builder::with_capacity(len);
205204

206205
for i in 0..len {
207206
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 {
218217
b.append_null();
219218
continue;
220219
}
221-
let next_bucket = (buckets + 1) as i32;
220+
let next_bucket = (buckets + 1) as i64;
222221
if $check_nan {
223222
if !x.is_finite() || !l.is_finite() || !h.is_finite() {
224223
b.append_null();
@@ -264,7 +263,7 @@ macro_rules! width_bucket_kernel_impl {
264263
b.append_null();
265264
continue;
266265
}
267-
let mut bucket = ((x - l) / width).floor() as i32 + 1;
266+
let mut bucket = ((x - l) / width).floor() as i64 + 1;
268267
if bucket < 1 {
269268
bucket = 1;
270269
}
@@ -306,9 +305,9 @@ pub(crate) fn width_bucket_interval_mdn_exact(
306305
lo: &IntervalMonthDayNanoArray,
307306
hi: &IntervalMonthDayNanoArray,
308307
n: &Int64Array,
309-
) -> Int32Array {
308+
) -> Int64Array {
310309
let len = v.len();
311-
let mut b = Int32Builder::with_capacity(len);
310+
let mut b = Int64Builder::with_capacity(len);
312311

313312
for i in 0..len {
314313
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(
320319
b.append_null();
321320
continue;
322321
}
323-
let next_bucket = (buckets + 1) as i32;
322+
let next_bucket = buckets + 1;
324323

325324
let x = v.value(i);
326325
let l = lo.value(i);
@@ -366,7 +365,7 @@ pub(crate) fn width_bucket_interval_mdn_exact(
366365
continue;
367366
}
368367

369-
let mut bucket = ((x_m - l_m) / width).floor() as i32 + 1;
368+
let mut bucket = ((x_m - l_m) / width).floor() as i64 + 1;
370369
if bucket < 1 {
371370
bucket = 1;
372371
}
@@ -417,7 +416,7 @@ pub(crate) fn width_bucket_interval_mdn_exact(
417416
continue;
418417
}
419418

420-
let mut bucket = ((x_f - l_f) / width).floor() as i32 + 1;
419+
let mut bucket = ((x_f - l_f) / width).floor() as i64 + 1;
421420
if bucket < 1 {
422421
bucket = 1;
423422
}
@@ -437,10 +436,11 @@ pub(crate) fn width_bucket_interval_mdn_exact(
437436
#[cfg(test)]
438437
mod tests {
439438
use super::*;
439+
use arrow::datatypes::Int64Type;
440440

441441
use arrow::array::{
442-
ArrayRef, DurationMicrosecondArray, Float64Array, Int32Array, Int64Array,
443-
IntervalYearMonthArray,
442+
ArrayRef, AsArray, DurationMicrosecondArray, Float64Array, Int32Array,
443+
Int64Array, IntervalYearMonthArray,
444444
};
445445
use arrow::datatypes::IntervalMonthDayNano;
446446

@@ -466,10 +466,6 @@ mod tests {
466466
Arc::new(IntervalYearMonthArray::from(vals.to_vec()))
467467
}
468468

469-
fn downcast_i32(arr: &ArrayRef) -> &Int32Array {
470-
arr.as_any().downcast_ref::<Int32Array>().unwrap()
471-
}
472-
473469
fn mdn_array(vals: &[(i32, i32, i64)]) -> Arc<IntervalMonthDayNanoArray> {
474470
let data: Vec<IntervalMonthDayNano> = vals
475471
.iter()
@@ -488,7 +484,7 @@ mod tests {
488484
let n = i64_array_all(5, 10);
489485

490486
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
491-
let out = downcast_i32(&out);
487+
let out = out.as_primitive::<Int64Type>();
492488
assert_eq!(out.values(), &[1, 2, 10, 0, 11]);
493489
}
494490

@@ -500,7 +496,7 @@ mod tests {
500496
let n = i64_array_all(5, 10);
501497

502498
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
503-
let out = downcast_i32(&out);
499+
let out = out.as_primitive::<Int64Type>();
504500

505501
assert_eq!(out.values(), &[1, 1, 11, 11, 0]);
506502
}
@@ -512,7 +508,7 @@ mod tests {
512508
let n = i64_array_all(3, 10);
513509

514510
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
515-
let out = downcast_i32(&out);
511+
let out = out.as_primitive::<Int64Type>();
516512
assert_eq!(out.values(), &[1, 10, 11]);
517513
}
518514

@@ -524,7 +520,7 @@ mod tests {
524520
let n = i64_array_all(3, 10);
525521

526522
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
527-
let out = downcast_i32(&out);
523+
let out = out.as_primitive::<Int64Type>();
528524
assert_eq!(out.values(), &[1, 11, 11]);
529525
}
530526

@@ -535,7 +531,7 @@ mod tests {
535531
let hi = f64_array(&[10.0, 10.0, 10.0]);
536532
let n = Arc::new(Int64Array::from(vec![0, -1, 10]));
537533
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
538-
let out = downcast_i32(&out);
534+
let out = out.as_primitive::<Int64Type>();
539535
assert!(out.is_null(0));
540536
assert!(out.is_null(1));
541537
assert_eq!(out.value(2), 10);
@@ -545,15 +541,15 @@ mod tests {
545541
let hi = f64_array(&[5.0]);
546542
let n = i64_array_all(1, 10);
547543
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
548-
let out = downcast_i32(&out);
544+
let out = out.as_primitive::<Int64Type>();
549545
assert!(out.is_null(0));
550546

551547
let v = f64_array_opt(&[Some(f64::NAN)]);
552548
let lo = f64_array(&[0.0]);
553549
let hi = f64_array(&[10.0]);
554550
let n = i64_array_all(1, 10);
555551
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
556-
let out = downcast_i32(&out);
552+
let out = out.as_primitive::<Int64Type>();
557553
assert!(out.is_null(0));
558554
}
559555

@@ -565,7 +561,7 @@ mod tests {
565561
let n = i64_array_all(4, 10);
566562

567563
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
568-
let out = downcast_i32(&out);
564+
let out = out.as_primitive::<Int64Type>();
569565
assert!(out.is_null(0));
570566
assert_eq!(out.value(1), 2);
571567
assert_eq!(out.value(2), 3);
@@ -576,7 +572,7 @@ mod tests {
576572
let hi = f64_array(&[10.0]);
577573
let n = i64_array_all(1, 10);
578574
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
579-
let out = downcast_i32(&out);
575+
let out = out.as_primitive::<Int64Type>();
580576
assert!(out.is_null(0));
581577
}
582578

@@ -590,7 +586,7 @@ mod tests {
590586
let n = i64_array_all(3, 2);
591587

592588
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
593-
let out = downcast_i32(&out);
589+
let out = out.as_primitive::<Int64Type>();
594590
assert_eq!(out.values(), &[2, 1, 0]);
595591
}
596592

@@ -601,7 +597,7 @@ mod tests {
601597
let hi = dur_us_array(&[1]);
602598
let n = i64_array_all(1, 10);
603599
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
604-
assert!(downcast_i32(&out).is_null(0));
600+
assert!(out.as_primitive::<Int64Type>().is_null(0));
605601
}
606602

607603
// --- Interval(YearMonth) ------------------------------------------------
@@ -614,7 +610,7 @@ mod tests {
614610
let n = i64_array_all(5, 12);
615611

616612
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
617-
let out = downcast_i32(&out);
613+
let out = out.as_primitive::<Int64Type>();
618614
assert_eq!(out.values(), &[1, 6, 12, 13, 13]);
619615
}
620616

@@ -626,7 +622,7 @@ mod tests {
626622
let n = i64_array_all(5, 12);
627623

628624
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
629-
let out = downcast_i32(&out);
625+
let out = out.as_primitive::<Int64Type>();
630626
assert_eq!(out.values(), &[2, 1, 13, 13, 0]);
631627
}
632628

@@ -640,7 +636,7 @@ mod tests {
640636
let n = i64_array_all(5, 12);
641637

642638
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
643-
let out = downcast_i32(&out);
639+
let out = out.as_primitive::<Int64Type>();
644640
assert_eq!(out.values(), &[1, 6, 12, 13, 13]);
645641
}
646642

@@ -652,7 +648,7 @@ mod tests {
652648
let n = i64_array_all(5, 12);
653649

654650
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
655-
let out = downcast_i32(&out);
651+
let out = out.as_primitive::<Int64Type>();
656652
// Mismo patrón que YM descendente
657653
assert_eq!(out.values(), &[2, 1, 13, 13, 0]);
658654
}
@@ -672,7 +668,7 @@ mod tests {
672668
let n = i64_array_all(6, 10);
673669

674670
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
675-
let out = downcast_i32(&out);
671+
let out = out.as_primitive::<Int64Type>();
676672
// x==hi -> n+1, x<lo -> 0, x>hi -> n+1
677673
assert_eq!(out.values(), &[1, 6, 10, 11, 0, 11]);
678674
}
@@ -685,7 +681,7 @@ mod tests {
685681
let n = i64_array_all(5, 10);
686682

687683
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
688-
let out = downcast_i32(&out);
684+
let out = out.as_primitive::<Int64Type>();
689685

690686
assert_eq!(out.values(), &[2, 1, 11, 11, 0]);
691687
}
@@ -697,7 +693,7 @@ mod tests {
697693
let n = i64_array_all(5, 10);
698694

699695
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
700-
let out = downcast_i32(&out);
696+
let out = out.as_primitive::<Int64Type>();
701697

702698
assert_eq!(out.values(), &[1, 1, 11, 11, 0]);
703699
}
@@ -710,7 +706,7 @@ mod tests {
710706
let n = i64_array_all(1, 4);
711707

712708
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
713-
let out = downcast_i32(&out);
709+
let out = out.as_primitive::<Int64Type>();
714710
assert!(out.is_null(0));
715711
}
716712

@@ -722,7 +718,7 @@ mod tests {
722718
let n = i64_array_all(1, 10);
723719

724720
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
725-
assert!(downcast_i32(&out).is_null(0));
721+
assert!(out.as_primitive::<Int64Type>().is_null(0));
726722
}
727723

728724
#[test]
@@ -733,7 +729,7 @@ mod tests {
733729
let n = Arc::new(Int64Array::from(vec![0])); // n <= 0
734730

735731
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
736-
assert!(downcast_i32(&out).is_null(0));
732+
assert!(out.as_primitive::<Int64Type>().is_null(0));
737733
}
738734

739735
#[test]
@@ -747,7 +743,7 @@ mod tests {
747743
let n = i64_array_all(2, 10);
748744

749745
let out = width_bucket_kern(&[v, lo, hi, n]).unwrap();
750-
let out = downcast_i32(&out);
746+
let out = out.as_primitive::<Int64Type>();
751747
assert!(out.is_null(0));
752748
assert_eq!(out.value(1), 6);
753749
}

0 commit comments

Comments
 (0)