[arrow-rs] branch master updated: Derive ArrowPrimitiveType for Decimal128Type and Decimal256Type (#2637) (#2833)

[tu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 37c867921 Derive ArrowPrimitiveType for Decimal128Type and Decimal256Type (#2637) (#2833)
37c867921 is described below

commit 37c867921c434974e908b22bafe0fc84f0de2ad0
Author: Raphael Taylor-Davies <17...@users.noreply.github.com>
AuthorDate: Thu Oct 6 20:48:13 2022 +0100

    Derive ArrowPrimitiveType for Decimal128Type and Decimal256Type (#2637) (#2833)
---
 arrow-array/src/array/primitive_array.rs | 44 +++++++++++++++++++++++++++++++-
 arrow-array/src/types.rs                 | 13 ++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
index e5f5cd481..3550e291c 100644
--- a/arrow-array/src/array/primitive_array.rs
+++ b/arrow-array/src/array/primitive_array.rs
@@ -22,7 +22,7 @@ use crate::temporal_conversions::{as_date, as_datetime, as_duration, as_time};
 use crate::trusted_len::trusted_len_unzip;
 use crate::types::*;
 use crate::{print_long_array, Array, ArrayAccessor};
-use arrow_buffer::{bit_util, ArrowNativeType, Buffer, MutableBuffer};
+use arrow_buffer::{bit_util, i256, ArrowNativeType, Buffer, MutableBuffer};
 use arrow_data::bit_iterator::try_for_each_valid_idx;
 use arrow_data::ArrayData;
 use arrow_schema::DataType;
@@ -609,6 +609,8 @@ def_from_for_primitive!(UInt64Type, u64);
 def_from_for_primitive!(Float16Type, f16);
 def_from_for_primitive!(Float32Type, f32);
 def_from_for_primitive!(Float64Type, f64);
+def_from_for_primitive!(Decimal128Type, i128);
+def_from_for_primitive!(Decimal256Type, i256);
 
 impl<T: ArrowPrimitiveType> From<Option<<T as ArrowPrimitiveType>::Native>>
     for NativeAdapter<T>
@@ -733,6 +735,8 @@ def_numeric_from_vec!(UInt32Type);
 def_numeric_from_vec!(UInt64Type);
 def_numeric_from_vec!(Float32Type);
 def_numeric_from_vec!(Float64Type);
+def_numeric_from_vec!(Decimal128Type);
+def_numeric_from_vec!(Decimal256Type);
 
 def_numeric_from_vec!(Date32Type);
 def_numeric_from_vec!(Date64Type);
@@ -1342,4 +1346,42 @@ mod tests {
 
         array.value(4);
     }
+
+    #[test]
+    fn test_decimal128() {
+        let values: Vec<_> = vec![0, 1, -1, i128::MIN, i128::MAX];
+        let array: PrimitiveArray<Decimal128Type> =
+            PrimitiveArray::from_iter(values.iter().copied());
+        assert_eq!(array.values(), &values);
+
+        let array: PrimitiveArray<Decimal128Type> =
+            PrimitiveArray::from_iter_values(values.iter().copied());
+        assert_eq!(array.values(), &values);
+
+        let array = PrimitiveArray::<Decimal128Type>::from(values.clone());
+        assert_eq!(array.values(), &values);
+
+        let array = PrimitiveArray::<Decimal128Type>::from(array.data().clone());
+        assert_eq!(array.values(), &values);
+    }
+
+    #[test]
+    fn test_decimal256() {
+        let values: Vec<_> =
+            vec![i256::ZERO, i256::ONE, i256::MINUS_ONE, i256::MIN, i256::MAX];
+
+        let array: PrimitiveArray<Decimal256Type> =
+            PrimitiveArray::from_iter(values.iter().copied());
+        assert_eq!(array.values(), &values);
+
+        let array: PrimitiveArray<Decimal256Type> =
+            PrimitiveArray::from_iter_values(values.iter().copied());
+        assert_eq!(array.values(), &values);
+
+        let array = PrimitiveArray::<Decimal256Type>::from(values.clone());
+        assert_eq!(array.values(), &values);
+
+        let array = PrimitiveArray::<Decimal256Type>::from(array.data().clone());
+        assert_eq!(array.values(), &values);
+    }
 }
diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs
index 581fdc767..9bd433692 100644
--- a/arrow-array/src/types.rs
+++ b/arrow-array/src/types.rs
@@ -19,6 +19,7 @@
 
 use crate::array::ArrowPrimitiveType;
 use crate::delta::shift_months;
+use arrow_buffer::i256;
 use arrow_data::decimal::{
     DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE, DECIMAL256_MAX_PRECISION,
     DECIMAL256_MAX_SCALE, DECIMAL_DEFAULT_SCALE,
@@ -515,6 +516,12 @@ impl DecimalType for Decimal128Type {
         DataType::Decimal128(DECIMAL128_MAX_PRECISION, DECIMAL_DEFAULT_SCALE);
 }
 
+impl ArrowPrimitiveType for Decimal128Type {
+    type Native = i128;
+
+    const DATA_TYPE: DataType = <Self as DecimalType>::DEFAULT_TYPE;
+}
+
 /// The decimal type for a Decimal256Array
 #[derive(Debug)]
 pub struct Decimal256Type {}
@@ -530,6 +537,12 @@ impl DecimalType for Decimal256Type {
         DataType::Decimal256(DECIMAL256_MAX_PRECISION, DECIMAL_DEFAULT_SCALE);
 }
 
+impl ArrowPrimitiveType for Decimal256Type {
+    type Native = i256;
+
+    const DATA_TYPE: DataType = <Self as DecimalType>::DEFAULT_TYPE;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;