For review

apache · Nov 9, 2022 · 1af7263 · 1af7263
1 parent 3e6d2bb
commit 1af7263
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 20 deletions.
diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
@@ -832,7 +832,7 @@ impl<T: DecimalType + ArrowPrimitiveType> PrimitiveArray<T> {
     /// 2. `scale` is larger than `T::MAX_SCALE`
     /// 3. `scale` is > `precision`
     pub fn with_precision_and_scale(
-        &self,
+        self,
         precision: u8,
         scale: u8,
     ) -> Result<Self, ArrowError>

diff --git a/arrow/src/row/fixed.rs b/arrow/src/row/fixed.rs
@@ -326,9 +326,20 @@ fn decode_fixed<T: FixedLengthEncoding + ToByteSlice>(
 pub fn decode_primitive<T: ArrowPrimitiveType>(
     rows: &mut [&[u8]],
     options: SortOptions,
+    data_type: &DataType,
 ) -> PrimitiveArray<T>
 where
     T::Native: FixedLengthEncoding + ToByteSlice,
 {
-    decode_fixed::<T::Native>(rows, T::DATA_TYPE, options).into()
+    let array_data: ArrayData = decode_fixed::<T::Native>(rows, T::DATA_TYPE, options);
+    match data_type {
+        DataType::Decimal128(_, _)
+        | DataType::Decimal256(_, _)
+        | DataType::Timestamp(_, _) => {
+            let data = array_data.into_builder().data_type(data_type.clone());
+
+            unsafe { data.build_unchecked().into() }
+        }
+        _ => array_data.into(),
+    }
 }
diff --git a/arrow/src/row/mod.rs b/arrow/src/row/mod.rs
@@ -610,8 +610,8 @@ fn encode_column(
 }
 
 macro_rules! decode_primitive_helper {
-    ($t:ty, $rows: ident, $options:ident) => {
-        Arc::new(decode_primitive::<$t>($rows, $options))
+    ($t:ty, $rows: ident, $options:ident, $data_type: ident) => {
+        Arc::new(decode_primitive::<$t>($rows, $options, $data_type))
     };
 }
 
@@ -626,8 +626,9 @@ unsafe fn decode_column(
     interner: Option<&OrderPreservingInterner>,
 ) -> Result<ArrayRef> {
     let options = field.options;
+    let data_type = &field.data_type;
     let array: ArrayRef = downcast_primitive! {
-        &field.data_type => (decode_primitive_helper, rows, options),
+        &field.data_type => (decode_primitive_helper, rows, options, data_type),
         DataType::Null => Arc::new(NullArray::new(rows.len())),
         DataType::Boolean => Arc::new(decode_bool(rows, options)),
         DataType::Binary => Arc::new(decode_binary::<i32>(rows, options)),
@@ -697,21 +698,6 @@ unsafe fn decode_column(
             )))
         }
     };
-    let array: ArrayRef = match &field.data_type {
-        DataType::Decimal128(p, s) => {
-            let d = as_primitive_array::<Decimal128Type>(&array)
-                .with_precision_and_scale(*p, *s)
-                .unwrap();
-            Arc::new(d)
-        }
-        DataType::Decimal256(p, s) => {
-            let d = as_primitive_array::<Decimal256Type>(&array)
-                .with_precision_and_scale(*p, *s)
-                .unwrap();
-            Arc::new(d)
-        }
-        _ => array,
-    };
     Ok(array)
 }