Derive ArrowPrimitiveType for Decimal128Type and Decimal256Type (apac…

…he#2637)
tustvold · Oct 6, 2022 · f1402ad · f1402ad
1 parent f8c4037
commit f1402ad
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
@@ -827,6 +827,7 @@ impl<T: ArrowPrimitiveType> From<ArrayData> for PrimitiveArray<T> {
 mod tests {
     use super::*;
     use crate::BooleanArray;
+    use arrow_buffer::i256;
 
     #[test]
     fn test_primitive_array_from_vec() {
@@ -1342,4 +1343,18 @@ mod tests {
 
         array.value(4);
     }
+
+    #[test]
+    fn test_decimal() {
+        let values: Vec<_> = vec![0, 1, -1, i128::MIN, i128::MAX];
+        let array: PrimitiveArray<Decimal128Type> =
+            PrimitiveArray::from_iter_values(values.iter().copied());
+        assert_eq!(array.values(), &values);
+
+        let values: Vec<_> =
+            vec![i256::ZERO, i256::ONE, i256::MINUS_ONE, i256::MIN, i256::MAX];
+        let array: PrimitiveArray<Decimal256Type> =
+            PrimitiveArray::from_iter_values(values.iter().copied());
+        assert_eq!(array.values(), &values);
+    }
 }
diff --git 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::*;