From 1c8ad26b4462d8225a4cc40fabe30fe794244e9a Mon Sep 17 00:00:00 2001 From: izveigor Date: Thu, 13 Apr 2023 19:22:39 +0300 Subject: [PATCH 1/2] feat: Support dyn_compare_scalar for Decimal256 --- arrow-ord/src/comparison.rs | 67 +++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/arrow-ord/src/comparison.rs b/arrow-ord/src/comparison.rs index 2927354da291..0444a0ae38af 100644 --- a/arrow-ord/src/comparison.rs +++ b/arrow-ord/src/comparison.rs @@ -26,6 +26,7 @@ use arrow_array::cast::*; use arrow_array::types::*; use arrow_array::*; +use arrow_buffer::i256; use arrow_buffer::{bit_util, BooleanBuffer, Buffer, MutableBuffer, NullBuffer}; use arrow_data::ArrayData; use arrow_schema::{ArrowError, DataType, IntervalUnit, TimeUnit}; @@ -517,6 +518,11 @@ macro_rules! dyn_compare_scalar { let left = as_primitive_array::($LEFT); $OP::(left, right) } + DataType::Decimal256(_, _) => { + let right = try_to_type!($RIGHT, to_i128)?; + let left = as_primitive_array::($LEFT); + $OP::(left, i256::from_i128(right)) + } DataType::Date32 => { let right = try_to_type!($RIGHT, to_i32)?; let left = as_primitive_array::($LEFT); @@ -6165,6 +6171,67 @@ mod tests { assert_eq!(e, r); } + #[test] + fn test_decimal256_scalar() { + let a = Decimal256Array::from_iter_values( + [1, 2, 3, 4, 5].into_iter().map(i256::from_i128), + ); + let b = i256::from_i128(3); + // array eq scalar + let e = BooleanArray::from( + vec![Some(false), Some(false), Some(true), Some(false), Some(false)], + ); + let r = eq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = eq_dyn_scalar(&a, b).unwrap(); + assert_eq!(e, r); + + // array neq scalar + let e = BooleanArray::from( + vec![Some(true), Some(true), Some(false), Some(true), Some(true)], + ); + let r = neq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = neq_dyn_scalar(&a, b).unwrap(); + assert_eq!(e, r); + + // array lt scalar + let e = BooleanArray::from( + vec![Some(true), Some(true), Some(false), Some(false), Some(false)], + ); + let r = lt_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = lt_dyn_scalar(&a, b).unwrap(); + assert_eq!(e, r); + + // array lt_eq scalar + let e = BooleanArray::from( + vec![Some(true), Some(true), Some(true), Some(false), Some(false)], + ); + let r = lt_eq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = lt_eq_dyn_scalar(&a, b).unwrap(); + assert_eq!(e, r); + + // array gt scalar + let e = BooleanArray::from( + vec![Some(false), Some(false), Some(false), Some(true), Some(true)], + ); + let r = gt_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = gt_dyn_scalar(&a, b).unwrap(); + assert_eq!(e, r); + + // array gt_eq scalar + let e = BooleanArray::from( + vec![Some(false), Some(false), Some(true), Some(true), Some(true)], + ); + let r = gt_eq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = gt_eq_dyn_scalar(&a, b).unwrap(); + assert_eq!(e, r); + } + #[test] #[cfg(not(feature = "simd"))] fn test_floating_zeros() { From f8232a3e2a777f0c1ccca5840892dde3e2deeccf Mon Sep 17 00:00:00 2001 From: izveigor Date: Thu, 13 Apr 2023 20:28:43 +0300 Subject: [PATCH 2/2] feat: add test_decimal256_scalar_i256 --- arrow-ord/src/comparison.rs | 63 ++++++++++++++++++++++++++++++++++++- 1 file changed, 62 insertions(+), 1 deletion(-) diff --git a/arrow-ord/src/comparison.rs b/arrow-ord/src/comparison.rs index 0444a0ae38af..b9ffddd5e1b4 100644 --- a/arrow-ord/src/comparison.rs +++ b/arrow-ord/src/comparison.rs @@ -6172,7 +6172,7 @@ mod tests { } #[test] - fn test_decimal256_scalar() { + fn test_decimal256_scalar_i128() { let a = Decimal256Array::from_iter_values( [1, 2, 3, 4, 5].into_iter().map(i256::from_i128), ); @@ -6232,6 +6232,67 @@ mod tests { assert_eq!(e, r); } + #[test] + fn test_decimal256_scalar_i256() { + let a = Decimal256Array::from_iter_values( + [1, 2, 3, 4, 5].into_iter().map(i256::from_i128), + ); + let b = i256::MAX; + // array eq scalar + let e = BooleanArray::from( + vec![Some(false), Some(false), Some(false), Some(false), Some(false)], + ); + let r = eq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = eq_dyn_scalar(&a, b).is_err(); + assert!(r); + + // array neq scalar + let e = BooleanArray::from( + vec![Some(true), Some(true), Some(true), Some(true), Some(true)], + ); + let r = neq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = neq_dyn_scalar(&a, b).is_err(); + assert!(r); + + // array lt scalar + let e = BooleanArray::from( + vec![Some(true), Some(true), Some(true), Some(true), Some(true)], + ); + let r = lt_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = lt_dyn_scalar(&a, b).is_err(); + assert!(r); + + // array lt_eq scalar + let e = BooleanArray::from( + vec![Some(true), Some(true), Some(true), Some(true), Some(true)], + ); + let r = lt_eq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = lt_eq_dyn_scalar(&a, b).is_err(); + assert!(r); + + // array gt scalar + let e = BooleanArray::from( + vec![Some(false), Some(false), Some(false), Some(false), Some(false)], + ); + let r = gt_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = gt_dyn_scalar(&a, b).is_err(); + assert!(r); + + // array gt_eq scalar + let e = BooleanArray::from( + vec![Some(false), Some(false), Some(false), Some(false), Some(false)], + ); + let r = gt_eq_scalar(&a, b).unwrap(); + assert_eq!(e, r); + let r = gt_eq_dyn_scalar(&a, b).is_err(); + assert!(r); + } + #[test] #[cfg(not(feature = "simd"))] fn test_floating_zeros() {