diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs index e5f5cd481af..3550e291c03 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 From::Native>> for NativeAdapter @@ -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 = + PrimitiveArray::from_iter(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array: PrimitiveArray = + PrimitiveArray::from_iter_values(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::from(values.clone()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::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 = + PrimitiveArray::from_iter(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array: PrimitiveArray = + PrimitiveArray::from_iter_values(values.iter().copied()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::from(values.clone()); + assert_eq!(array.values(), &values); + + let array = PrimitiveArray::::from(array.data().clone()); + assert_eq!(array.values(), &values); + } } diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs index 581fdc767c2..9bd43369258 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 = ::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 = ::DEFAULT_TYPE; +} + #[cfg(test)] mod tests { use super::*;