ARROW-10818: [Rust] Implement DecimalType

apache · Jan 17, 2021 · 7e2fca8 · 7e2fca8
1 parent e73f205
commit 7e2fca8
Show file tree

Hide file tree

Showing 6 changed files with 416 additions and 15 deletions.
diff --git a/rust/arrow/Cargo.toml b/rust/arrow/Cargo.toml
@@ -51,6 +51,7 @@ flatbuffers = "^0.8"
 hex = "0.4"
 prettytable-rs = { version = "0.8.0", optional = true }
 lexical-core = "^0.7"
+num-bigint = "0.3"
 
 [features]
 default = []

diff --git a/rust/arrow/src/array/array_binary.rs b/rust/arrow/src/array/array_binary.rs
@@ -29,7 +29,7 @@ use super::{
 };
 use crate::util::bit_util;
 use crate::{buffer::Buffer, datatypes::ToByteSlice};
-use crate::{buffer::MutableBuffer, datatypes::DataType};
+use crate::{buffer::MutableBuffer, datatypes::DataType, datatypes::DecimalType};
 
 /// Like OffsetSizeTrait, but specialized for Binary
 // This allow us to expose a constant datatype for the GenericBinaryArray
@@ -485,9 +485,9 @@ pub struct DecimalArray {
 }
 
 impl DecimalArray {
-    /// Returns the element at index `i` as i128.
-    pub fn value(&self, i: usize) -> i128 {
+    pub fn value(&self, i: usize) -> DecimalType {
         assert!(i < self.data.len(), "DecimalArray out of bounds access");
+
         let offset = i.checked_add(self.data.offset()).unwrap();
         let raw_val = unsafe {
             let pos = self.value_offset_at(offset);
@@ -496,11 +496,8 @@ impl DecimalArray {
                 (self.value_offset_at(offset + 1) - pos) as usize,
             )
         };
-        let as_array = raw_val.try_into();
-        match as_array {
-            Ok(v) if raw_val.len() == 16 => i128::from_le_bytes(v),
-            _ => panic!("DecimalArray elements are not 128bit integers."),
-        }
+
+        DecimalType::from_bytes_with_scale(raw_val, self.scale)
     }
 
     /// Returns the offset for the element at index `i`.
@@ -987,13 +984,20 @@ mod tests {
             192, 219, 180, 17, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 64, 36, 75, 238, 253,
             255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
         ];
-        let array_data = ArrayData::builder(DataType::Decimal(23, 6))
+        let array_data = ArrayData::builder(DataType::Decimal(23, 0))
             .len(2)
             .add_buffer(Buffer::from(&values[..]))
             .build();
         let decimal_array = DecimalArray::from(array_data);
-        assert_eq!(8_887_000_000, decimal_array.value(0));
-        assert_eq!(-8_887_000_000, decimal_array.value(1));
+
+        assert_eq!(
+            DecimalType::from(8_887_000_000_i128),
+            decimal_array.value(0)
+        );
+        assert_eq!(
+            DecimalType::from(-8_887_000_000_i128),
+            decimal_array.value(1)
+        );
         assert_eq!(16, decimal_array.value_length());
     }
 
@@ -1009,7 +1013,7 @@ mod tests {
             .build();
         let arr = DecimalArray::from(array_data);
         assert_eq!(
-            "DecimalArray<23, 6>\n[\n  8887000000,\n  -8887000000,\n]",
+            "DecimalArray<23, 6>\n[\n  Decimal<6>(\"8887.000000\"),\n  Decimal<6>(\"-8887.000000\"),\n]",
             format!("{:?}", arr)
         );
     }

diff --git a/rust/arrow/src/array/equal_json.rs b/rust/arrow/src/array/equal_json.rs
@@ -345,7 +345,7 @@ impl JsonEqual for DecimalArray {
             JString(s) => {
                 self.is_valid(i)
                     && (s
-                        .parse::<i128>()
+                        .parse::<DecimalType>()
                         .map_or_else(|_| false, |v| v == self.value(i)))
             }
             JNull => self.is_null(i),