Add Decimal128 and Decimal256 to downcast_primitive (#3056)

* Add Decimal128 and Decimal256 to downcast_primitive * Add doc test * Fix test * For review
apache · Nov 9, 2022 · f596209 · f596209
1 parent e4e15f8
commit f596209
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 21 deletions.
diff --git a/arrow-array/src/cast.rs b/arrow-array/src/cast.rs
@@ -225,6 +225,8 @@ macro_rules! downcast_temporal_array {
 /// assert_eq!(primitive_size(&DataType::Int32), 4);
 /// assert_eq!(primitive_size(&DataType::Int64), 8);
 /// assert_eq!(primitive_size(&DataType::Float16), 2);
+/// assert_eq!(primitive_size(&DataType::Decimal128(38, 10)), 16);
+/// assert_eq!(primitive_size(&DataType::Decimal256(76, 20)), 32);
 /// ```
 ///
 /// [`DataType`]: arrow_schema::DataType
@@ -242,6 +244,12 @@ macro_rules! downcast_primitive {
             $crate::repeat_pat!(arrow_schema::DataType::Float64, $($data_type),+) => {
                 $m!($crate::types::Float64Type $(, $args)*)
             }
+            $crate::repeat_pat!(arrow_schema::DataType::Decimal128(_, _), $($data_type),+) => {
+                $m!($crate::types::Decimal128Type $(, $args)*)
+            }
+            $crate::repeat_pat!(arrow_schema::DataType::Decimal256(_, _), $($data_type),+) => {
+                $m!($crate::types::Decimal256Type $(, $args)*)
+            }
             $crate::repeat_pat!(arrow_schema::DataType::Interval(arrow_schema::IntervalUnit::YearMonth), $($data_type),+) => {
                 $m!($crate::types::IntervalYearMonthType $(, $args)*)
             }

diff --git a/arrow/src/row/mod.rs b/arrow/src/row/mod.rs
@@ -102,7 +102,6 @@ use std::sync::Arc;
 
 use arrow_array::cast::*;
 use arrow_array::*;
-use arrow_buffer::i256;
 
 use crate::compute::SortOptions;
 use crate::datatypes::*;
@@ -504,8 +503,6 @@ fn new_empty_rows(
             array => lengths.iter_mut().for_each(|x| *x += fixed::encoded_len(array)),
             DataType::Null => {},
             DataType::Boolean => lengths.iter_mut().for_each(|x| *x += bool::ENCODED_LEN),
-            DataType::Decimal128(_, _) => lengths.iter_mut().for_each(|x| *x += i128::ENCODED_LEN),
-            DataType::Decimal256(_, _) => lengths.iter_mut().for_each(|x| *x += i256::ENCODED_LEN),
             DataType::Binary => as_generic_binary_array::<i32>(array)
                 .iter()
                 .zip(lengths.iter_mut())
@@ -586,22 +583,6 @@ fn encode_column(
         column => fixed::encode(out, column, opts),
         DataType::Null => {}
         DataType::Boolean => fixed::encode(out, as_boolean_array(column), opts),
-        DataType::Decimal128(_, _) => {
-            let column = column
-                .as_any()
-                .downcast_ref::<Decimal128Array>()
-                .unwrap();
-
-            fixed::encode(out, column, opts)
-        },
-        DataType::Decimal256(_, _) => {
-            let column = column
-                .as_any()
-                .downcast_ref::<Decimal256Array>()
-                .unwrap();
-
-            fixed::encode(out, column, opts)
-        },
         DataType::Binary => {
             variable::encode(out, as_generic_binary_array::<i32>(column).iter(), opts)
         }
@@ -654,8 +635,6 @@ unsafe fn decode_column(
         DataType::LargeBinary => Arc::new(decode_binary::<i64>(rows, options)),
         DataType::Utf8 => Arc::new(decode_string::<i32>(rows, options)),
         DataType::LargeUtf8 => Arc::new(decode_string::<i64>(rows, options)),
-        DataType::Decimal128(_, _) => decode_primitive_helper!(Decimal128Type, rows, data_type, options),
-        DataType::Decimal256(_, _) => decode_primitive_helper!(Decimal256Type, rows, data_type, options),
         DataType::Dictionary(k, v) => match k.as_ref() {
             DataType::Int8 => Arc::new(decode_dictionary::<Int8Type>(
                 interner.unwrap(),