From 1ac1d0a1dec37a8cb6c1393622b38a2132f8585a Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Tue, 8 Mar 2022 21:14:25 +0000 Subject: [PATCH] Added support for min/max for decimal --- src/compute/aggregate/min_max.rs | 2 + src/compute/aggregate/simd/mod.rs | 60 +++++++++++++++++++++++++++ src/compute/aggregate/simd/native.rs | 54 +----------------------- tests/it/compute/aggregate/min_max.rs | 7 ++++ 4 files changed, 70 insertions(+), 53 deletions(-) diff --git a/src/compute/aggregate/min_max.rs b/src/compute/aggregate/min_max.rs index fe89bf41537..8b78614a1a7 100644 --- a/src/compute/aggregate/min_max.rs +++ b/src/compute/aggregate/min_max.rs @@ -394,6 +394,7 @@ pub fn max(array: &dyn Array) -> Result> { DataType::Float16 => unreachable!(), DataType::Float32 => dyn_primitive!(f32, array, max_primitive), DataType::Float64 => dyn_primitive!(f64, array, max_primitive), + DataType::Decimal(_, _) => dyn_primitive!(i128, array, max_primitive), DataType::Utf8 => dyn_generic!(Utf8Array, Utf8Scalar, array, max_string), DataType::LargeUtf8 => dyn_generic!(Utf8Array, Utf8Scalar, array, max_string), DataType::Binary => dyn_generic!(BinaryArray, BinaryScalar, array, max_binary), @@ -435,6 +436,7 @@ pub fn min(array: &dyn Array) -> Result> { DataType::Float16 => unreachable!(), DataType::Float32 => dyn_primitive!(f32, array, min_primitive), DataType::Float64 => dyn_primitive!(f64, array, min_primitive), + DataType::Decimal(_, _) => dyn_primitive!(i128, array, min_primitive), DataType::Utf8 => dyn_generic!(Utf8Array, Utf8Scalar, array, min_string), DataType::LargeUtf8 => dyn_generic!(Utf8Array, Utf8Scalar, array, min_string), DataType::Binary => dyn_generic!(BinaryArray, BinaryScalar, array, min_binary), diff --git a/src/compute/aggregate/simd/mod.rs b/src/compute/aggregate/simd/mod.rs index ac0cbc942de..f8cbb07e711 100644 --- a/src/compute/aggregate/simd/mod.rs +++ b/src/compute/aggregate/simd/mod.rs @@ -1,3 +1,63 @@ +use super::SimdOrd; +use crate::types::simd::{i128x8, NativeSimd}; + +macro_rules! simd_ord_int { + ($simd:tt, $type:ty) => { + impl SimdOrd<$type> for $simd { + const MIN: $type = <$type>::MIN; + const MAX: $type = <$type>::MAX; + + #[inline] + fn max_element(self) -> $type { + self.0.iter().copied().fold(Self::MIN, <$type>::max) + } + + #[inline] + fn min_element(self) -> $type { + self.0.iter().copied().fold(Self::MAX, <$type>::min) + } + + #[inline] + fn max_lane(self, x: Self) -> Self { + let mut result = <$simd>::default(); + result + .0 + .iter_mut() + .zip(self.0.iter()) + .zip(x.0.iter()) + .for_each(|((a, b), c)| *a = (*b).max(*c)); + result + } + + #[inline] + fn min_lane(self, x: Self) -> Self { + let mut result = <$simd>::default(); + result + .0 + .iter_mut() + .zip(self.0.iter()) + .zip(x.0.iter()) + .for_each(|((a, b), c)| *a = (*b).min(*c)); + result + } + + #[inline] + fn new_min() -> Self { + Self([Self::MAX; <$simd>::LANES]) + } + + #[inline] + fn new_max() -> Self { + Self([Self::MIN; <$simd>::LANES]) + } + } + }; +} + +pub(super) use simd_ord_int; + +simd_ord_int!(i128x8, i128); + #[cfg(not(feature = "simd"))] mod native; #[cfg(not(feature = "simd"))] diff --git a/src/compute/aggregate/simd/native.rs b/src/compute/aggregate/simd/native.rs index 3304d5ea2b0..1da285afe99 100644 --- a/src/compute/aggregate/simd/native.rs +++ b/src/compute/aggregate/simd/native.rs @@ -4,6 +4,7 @@ use crate::types::simd::*; use super::super::min_max::SimdOrd; use super::super::sum::Sum; +use super::simd_ord_int; macro_rules! simd_add { ($simd:tt, $type:ty, $lanes:expr, $add:tt) => { @@ -53,59 +54,6 @@ simd_add!(i64x8, i64, 8, wrapping_add); simd_add!(f32x16, f32, 16, add); simd_add!(f64x8, f64, 8, add); -macro_rules! simd_ord_int { - ($simd:tt, $type:ty) => { - impl SimdOrd<$type> for $simd { - const MIN: $type = <$type>::MIN; - const MAX: $type = <$type>::MAX; - - #[inline] - fn max_element(self) -> $type { - self.0.iter().copied().fold(Self::MIN, <$type>::max) - } - - #[inline] - fn min_element(self) -> $type { - self.0.iter().copied().fold(Self::MAX, <$type>::min) - } - - #[inline] - fn max_lane(self, x: Self) -> Self { - let mut result = <$simd>::default(); - result - .0 - .iter_mut() - .zip(self.0.iter()) - .zip(x.0.iter()) - .for_each(|((a, b), c)| *a = (*b).max(*c)); - result - } - - #[inline] - fn min_lane(self, x: Self) -> Self { - let mut result = <$simd>::default(); - result - .0 - .iter_mut() - .zip(self.0.iter()) - .zip(x.0.iter()) - .for_each(|((a, b), c)| *a = (*b).min(*c)); - result - } - - #[inline] - fn new_min() -> Self { - Self([Self::MAX; <$simd>::LANES]) - } - - #[inline] - fn new_max() -> Self { - Self([Self::MIN; <$simd>::LANES]) - } - } - }; -} - macro_rules! simd_ord_float { ($simd:tt, $type:ty) => { impl SimdOrd<$type> for $simd { diff --git a/tests/it/compute/aggregate/min_max.rs b/tests/it/compute/aggregate/min_max.rs index babf5b9f539..224a245ae79 100644 --- a/tests/it/compute/aggregate/min_max.rs +++ b/tests/it/compute/aggregate/min_max.rs @@ -25,6 +25,13 @@ fn test_primitive_min_max_1() { assert_eq!(Some(5), max_primitive(&a)); } +#[test] +fn decimal() { + let a = Int128Array::from(&[None, None, Some(5), Some(2)]); + assert_eq!(Some(2), min_primitive(&a)); + assert_eq!(Some(5), max_primitive(&a)); +} + #[test] fn min_max_f32() { let a = Float32Array::from(&[None, None, Some(5.0), Some(2.0)]);