fix: Decimal ToArrow conversion (#3367)

Author: a10y

.github/workflows/labels.yml

@@ -13,6 +13,7 @@ jobs:
           contains(github.event.pull_request.labels.*.name, 'chore') == false &&
           contains(github.event.pull_request.labels.*.name, 'bug') == false &&
           contains(github.event.pull_request.labels.*.name, 'feature') == false &&
+          contains(github.event.pull_request.labels.*.name, 'fix') == false &&
           contains(github.event.pull_request.labels.*.name, 'performance') == false &&
           contains(github.event.pull_request.labels.*.name, 'break') == false &&
           contains(github.event.pull_request.labels.*.name, 'wire-break') == false

encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs

@@ -1,8 +1,6 @@
 mod compute;
 mod serde;
 
-use std::iter;
-
 use itertools::Itertools;
 use vortex_array::arrays::DecimalArray;
 use vortex_array::stats::{ArrayStats, StatsSetRef};
@@ -82,27 +80,6 @@ impl DecimalBytePartsArray {
             vortex_bail!("decimal bytes parts 2nd to 4th must be non-nullable u64 primitive typed")
         }
 
-        let primitive_bit_width = iter::once(&msp)
-            .chain(&lower_parts)
-            .map(|a| {
-                PType::try_from(a.dtype())
-                    .vortex_expect("already checked")
-                    .bit_width()
-            })
-            .sum();
-
-        if decimal_dtype.required_bit_width() > primitive_bit_width {
-            vortex_bail!(
-                "cannot represent a decimal {decimal_dtype} as primitive parts {:?}, decimal bit width {}, primitive bit width {}",
-                iter::once(&msp)
-                    .chain(&lower_parts)
-                    .map(|a| a.dtype())
-                    .collect_vec(),
-                decimal_dtype.required_bit_width(),
-                primitive_bit_width
-            )
-        }
-
         let nullable = msp.dtype().nullability();
         Ok(Self {
             msp,

java/vortex-jni/src/main/java/dev/vortex/api/Array.java

@@ -15,6 +15,7 @@
  */
 package dev.vortex.api;
 
+import java.math.BigDecimal;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.VectorSchemaRoot;
 
@@ -52,6 +53,8 @@ public interface Array extends AutoCloseable {
 
     double getDouble(int index);
 
+    BigDecimal getBigDecimal(int index);
+
     String getUTF8(int index);
 
     void getUTF8_ptr_len(int index, long[] ptr, int[] len);

java/vortex-jni/src/main/java/dev/vortex/jni/JNIArray.java

@@ -18,6 +18,7 @@
 import com.google.common.base.Preconditions;
 import dev.vortex.api.Array;
 import dev.vortex.api.DType;
+import java.math.BigDecimal;
 import java.util.OptionalLong;
 import org.apache.arrow.c.ArrowArray;
 import org.apache.arrow.c.ArrowSchema;
@@ -130,6 +131,11 @@ public double getDouble(int index) {
         return NativeArrayMethods.getDouble(pointer.getAsLong(), index);
     }
 
+    @Override
+    public BigDecimal getBigDecimal(int index) {
+        return NativeArrayMethods.getBigDecimal(pointer.getAsLong(), index);
+    }
+
     @Override
     public String getUTF8(int index) {
         return NativeArrayMethods.getUTF8(pointer.getAsLong(), index);

java/vortex-jni/src/main/java/dev/vortex/jni/NativeArrayMethods.java

@@ -15,6 +15,8 @@
  */
 package dev.vortex.jni;
 
+import java.math.BigDecimal;
+
 public final class NativeArrayMethods {
     static {
         NativeLoader.loadJni();
@@ -58,6 +60,8 @@ private NativeArrayMethods() {}
 
     public static native double getDouble(long pointer, int index);
 
+    public static native BigDecimal getBigDecimal(long pointer, int index);
+
     public static native String getUTF8(long pointer, int index);
 
     /**

vortex-array/src/arrays/constant/canonical.rs

@@ -11,7 +11,7 @@ use crate::arrays::constant::ConstantArray;
 use crate::arrays::primitive::PrimitiveArray;
 use crate::arrays::{
     BinaryView, BoolArray, ConstantVTable, DecimalArray, ExtensionArray, ListArray, NullArray,
-    StructArray, VarBinViewArray, precision_to_storage_size,
+    StructArray, VarBinViewArray, smallest_storage_type,
 };
 use crate::builders::{ArrayBuilderExt, builder_with_capacity};
 use crate::validity::Validity;
@@ -57,7 +57,7 @@ impl CanonicalVTable<ConstantVTable> for ConstantVTable {
                 })
             }
             DType::Decimal(decimal_type, ..) => {
-                let size = precision_to_storage_size(decimal_type);
+                let size = smallest_storage_type(decimal_type);
                 let decimal = scalar.as_decimal();
                 let Some(value) = decimal.decimal_value() else {
                     let all_null = match_each_decimal_value_type!(size, |$D| {

vortex-array/src/arrays/decimal/mod.rs

@@ -44,8 +44,8 @@ impl VTable for DecimalVTable {
 #[derive(Clone, Debug)]
 pub struct DecimalEncoding;
 
-/// Maps a decimal precision into the small type that can represent it.
-pub fn precision_to_storage_size(decimal_dtype: &DecimalDType) -> DecimalValueType {
+/// Maps a decimal precision into the smallest type that can represent it.
+pub fn smallest_storage_type(decimal_dtype: &DecimalDType) -> DecimalValueType {
     match decimal_dtype.precision() {
         1..=2 => DecimalValueType::I8,
         3..=4 => DecimalValueType::I16,

vortex-array/src/arrow/compute/to_arrow/canonical.rs

@@ -14,10 +14,10 @@ use arrow_array::{
 use arrow_buffer::{ScalarBuffer, i256};
 use arrow_schema::{DataType, Field, FieldRef, Fields};
 use itertools::Itertools;
-use num_traits::AsPrimitive;
+use num_traits::{AsPrimitive, ToPrimitive};
 use vortex_buffer::Buffer;
 use vortex_dtype::{DType, NativePType, PType};
-use vortex_error::{VortexExpect, VortexResult, vortex_bail};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_scalar::DecimalValueType;
 
 use crate::arrays::{
@@ -104,8 +104,34 @@ impl Kernel for ToArrowCanonical {
             {
                 to_arrow_primitive::<Float64Type>(array)
             }
-            (Canonical::Decimal(array), DataType::Decimal128(..)) => to_arrow_decimal128(array),
-            (Canonical::Decimal(array), DataType::Decimal256(..)) => to_arrow_decimal256(array),
+            (Canonical::Decimal(array), DataType::Decimal128(precision, scale)) => {
+                if array.decimal_dtype().precision() != *precision
+                    || array.decimal_dtype().scale() != *scale
+                {
+                    vortex_bail!(
+                        "ToArrowCanonical: target precision/scale {}/{} does not match array precision/scale {}/{}",
+                        precision,
+                        scale,
+                        array.decimal_dtype().precision(),
+                        array.decimal_dtype().scale()
+                    );
+                }
+                to_arrow_decimal128(array)
+            }
+            (Canonical::Decimal(array), DataType::Decimal256(precision, scale)) => {
+                if array.decimal_dtype().precision() != *precision
+                    || array.decimal_dtype().scale() != *scale
+                {
+                    vortex_bail!(
+                        "ToArrowCanonical: target precision/scale {}/{} does not match array precision/scale {}/{}",
+                        precision,
+                        scale,
+                        array.decimal_dtype().precision(),
+                        array.decimal_dtype().scale()
+                    );
+                }
+                to_arrow_decimal256(array)
+            }
             (Canonical::Struct(array), DataType::Struct(fields)) => {
                 to_arrow_struct(array, fields.as_ref())
             }
@@ -188,9 +214,14 @@ fn to_arrow_decimal128(array: DecimalArray) -> VortexResult<ArrowArrayRef> {
         DecimalValueType::I32 => array.buffer::<i32>().into_iter().map(|x| x.as_()).collect(),
         DecimalValueType::I64 => array.buffer::<i64>().into_iter().map(|x| x.as_()).collect(),
         DecimalValueType::I128 => array.buffer::<i128>(),
-        DecimalValueType::I256 => {
-            vortex_bail!("i256 decimals cannot be converted to Arrow i128 decimal")
-        }
+        DecimalValueType::I256 => array
+            .buffer::<vortex_scalar::i256>()
+            .into_iter()
+            .map(|x| {
+                x.to_i128()
+                    .ok_or_else(|| vortex_err!("i256 to i128 narrowing cannot be done safely"))
+            })
+            .try_collect()?,
         _ => vortex_bail!("unknown value type {:?}", array.values_type()),
     };
     Ok(Arc::new(
@@ -206,10 +237,14 @@ fn to_arrow_decimal256(array: DecimalArray) -> VortexResult<ArrowArrayRef> {
     let null_buffer = array.validity_mask()?.to_null_buffer();
     let buffer: Buffer<i256> = match array.values_type() {
         DecimalValueType::I8 => array.buffer::<i8>().into_iter().map(|x| x.as_()).collect(),
-        DecimalValueType::I16 => array.buffer::<i8>().into_iter().map(|x| x.as_()).collect(),
-        DecimalValueType::I32 => array.buffer::<i8>().into_iter().map(|x| x.as_()).collect(),
-        DecimalValueType::I64 => array.buffer::<i8>().into_iter().map(|x| x.as_()).collect(),
-        DecimalValueType::I128 => array.buffer::<i8>().into_iter().map(|x| x.as_()).collect(),
+        DecimalValueType::I16 => array.buffer::<i16>().into_iter().map(|x| x.as_()).collect(),
+        DecimalValueType::I32 => array.buffer::<i32>().into_iter().map(|x| x.as_()).collect(),
+        DecimalValueType::I64 => array.buffer::<i64>().into_iter().map(|x| x.as_()).collect(),
+        DecimalValueType::I128 => array
+            .buffer::<i128>()
+            .into_iter()
+            .map(|x| vortex_scalar::i256::from_i128(x).into())
+            .collect(),
         DecimalValueType::I256 => Buffer::<i256>::from_byte_buffer(array.byte_buffer()),
         _ => vortex_bail!("unknown type {:?}", array.values_type()),
     };
@@ -334,15 +369,19 @@ where
 
 #[cfg(test)]
 mod tests {
-    use arrow_array::Decimal128Array;
+    use arrow_array::{Array, Decimal128Array, Decimal256Array};
+    use arrow_buffer::i256;
     use arrow_schema::{DataType, Field};
+    use rstest::rstest;
     use vortex_buffer::buffer;
     use vortex_dtype::{DecimalDType, FieldNames};
+    use vortex_scalar::NativeDecimalType;
 
     use crate::IntoArray;
     use crate::arrays::{DecimalArray, PrimitiveArray, StructArray};
     use crate::arrow::IntoArrowArray;
     use crate::arrow::compute::to_arrow;
+    use crate::builders::{ArrayBuilder, DecimalBuilder};
     use crate::validity::Validity;
 
     #[test]
@@ -398,4 +437,54 @@ mod tests {
 
         assert!(struct_a.into_array().into_arrow(&arrow_dt).is_err());
     }
+
+    #[rstest]
+    #[case(0i8)]
+    #[case(0i16)]
+    #[case(0i32)]
+    #[case(0i64)]
+    #[case(0i128)]
+    #[case(vortex_scalar::i256::ZERO)]
+    fn to_arrow_decimal128<T: NativeDecimalType>(#[case] _decimal_type: T) {
+        let mut decimal = DecimalBuilder::new::<T>(2, 1, false.into());
+        decimal.append_value(10);
+        decimal.append_value(11);
+        decimal.append_value(12);
+
+        let decimal = decimal.finish();
+
+        let arrow_array = decimal.into_arrow(&DataType::Decimal128(2, 1)).unwrap();
+        let arrow_decimal = arrow_array
+            .as_any()
+            .downcast_ref::<Decimal128Array>()
+            .unwrap();
+        assert_eq!(arrow_decimal.value(0), 10);
+        assert_eq!(arrow_decimal.value(1), 11);
+        assert_eq!(arrow_decimal.value(2), 12);
+    }
+
+    #[rstest]
+    #[case(0i8)]
+    #[case(0i16)]
+    #[case(0i32)]
+    #[case(0i64)]
+    #[case(0i128)]
+    #[case(vortex_scalar::i256::ZERO)]
+    fn to_arrow_decimal256<T: NativeDecimalType>(#[case] _decimal_type: T) {
+        let mut decimal = DecimalBuilder::new::<T>(2, 1, false.into());
+        decimal.append_value(10);
+        decimal.append_value(11);
+        decimal.append_value(12);
+
+        let decimal = decimal.finish();
+
+        let arrow_array = decimal.into_arrow(&DataType::Decimal256(2, 1)).unwrap();
+        let arrow_decimal = arrow_array
+            .as_any()
+            .downcast_ref::<Decimal256Array>()
+            .unwrap();
+        assert_eq!(arrow_decimal.value(0), i256::from_i128(10));
+        assert_eq!(arrow_decimal.value(1), i256::from_i128(11));
+        assert_eq!(arrow_decimal.value(2), i256::from_i128(12));
+    }
 }

vortex-array/src/arrow/convert.rs

@@ -98,24 +98,24 @@ impl_from_arrow_primitive!(Float32Type);
 impl_from_arrow_primitive!(Float64Type);
 
 impl FromArrowArray<&ArrowPrimitiveArray<Decimal128Type>> for ArrayRef {
-    fn from_arrow(array: &ArrowPrimitiveArray<Decimal128Type>, _nullable: bool) -> Self {
+    fn from_arrow(array: &ArrowPrimitiveArray<Decimal128Type>, nullable: bool) -> Self {
         let decimal_type = DecimalDType::new(array.precision(), array.scale());
         let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
-        let validity = nulls(array.nulls(), false);
+        let validity = nulls(array.nulls(), nullable);
         DecimalArray::new(buffer, decimal_type, validity).into_array()
     }
 }
 
 impl FromArrowArray<&ArrowPrimitiveArray<Decimal256Type>> for ArrayRef {
-    fn from_arrow(array: &ArrowPrimitiveArray<Decimal256Type>, _nullable: bool) -> Self {
+    fn from_arrow(array: &ArrowPrimitiveArray<Decimal256Type>, nullable: bool) -> Self {
         let decimal_type = DecimalDType::new(array.precision(), array.scale());
         let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
         // SAFETY: Our i256 implementation has the same bit-pattern representation of the
         //  arrow_buffer::i256 type. It is safe to treat values held inside the buffer as values
         //  of either type.
         let buffer =
             unsafe { std::mem::transmute::<Buffer<arrow_buffer::i256>, Buffer<i256>>(buffer) };
-        let validity = nulls(array.nulls(), false);
+        let validity = nulls(array.nulls(), nullable);
         DecimalArray::new(buffer, decimal_type, validity).into_array()
     }
 }