vortex_array/arrays/decimal/compute/
to_arrow.rs

1use std::mem;
2use std::sync::Arc;
3
4use arrow_array::{ArrayRef as ArrowArrayRef, ArrayRef, Decimal128Array, Decimal256Array};
5use arrow_schema::DataType;
6use num_traits::AsPrimitive;
7use vortex_buffer::Buffer;
8use vortex_error::{VortexResult, vortex_bail};
9use vortex_scalar::i256;
10
11use crate::Array;
12use crate::arrays::decimal::serde::DecimalValueType;
13use crate::arrays::{DecimalArray, DecimalEncoding, NativeDecimalType};
14use crate::compute::ToArrowFn;
15
16impl ToArrowFn<&DecimalArray> for DecimalEncoding {
17    fn to_arrow(
18        &self,
19        array: &DecimalArray,
20        data_type: &DataType,
21    ) -> VortexResult<Option<ArrowArrayRef>> {
22        let precision = array.decimal_dtype().precision();
23        let scale = array.decimal_dtype().scale();
24        let nulls = array.validity_mask()?.to_null_buffer();
25
26        match array.values_type {
27            DecimalValueType::I8 => {
28                decimal_into_array_i128_decimal::<i8>(array, data_type, convert_buffer).map(Some)
29            }
30            DecimalValueType::I16 => {
31                decimal_into_array_i128_decimal::<i16>(array, data_type, convert_buffer).map(Some)
32            }
33            DecimalValueType::I32 => {
34                decimal_into_array_i128_decimal::<i32>(array, data_type, convert_buffer).map(Some)
35            }
36            DecimalValueType::I64 => {
37                decimal_into_array_i128_decimal::<i64>(array, data_type, convert_buffer).map(Some)
38            }
39            DecimalValueType::I128 => {
40                decimal_into_array_i128_decimal::<i128>(array, data_type, |x| x).map(Some)
41            }
42            DecimalValueType::I256 => {
43                let DataType::Decimal256(p, s) = data_type else {
44                    vortex_bail!(
45                        "Target Arrow type does not match Decimal source: {:?} ≠ {:?}",
46                        data_type,
47                        array.decimal_dtype()
48                    );
49                };
50                if *p != precision || *s != scale {
51                    vortex_bail!(
52                        "Decimal128: precision {} and scale {} do not match expected ({}, {})",
53                        precision,
54                        scale,
55                        p,
56                        s
57                    );
58                }
59
60                let buffer_i256 = array.buffer::<i256>();
61                // SAFETY: vortex_scalar::i256 and arrow_buffer::i256 have same bits
62                let buffer_i256: Buffer<arrow_buffer::i256> =
63                    unsafe { mem::transmute(buffer_i256) };
64
65                Ok(Some(Arc::new(
66                    Decimal256Array::new(buffer_i256.into_arrow_scalar_buffer(), nulls)
67                        .with_precision_and_scale(precision, scale)?,
68                )))
69            }
70        }
71    }
72}
73
74fn decimal_into_array_i128_decimal<T: NativeDecimalType + AsPrimitive<i128>>(
75    array: &DecimalArray,
76    data_type: &DataType,
77    convert: impl Fn(Buffer<T>) -> Buffer<i128>,
78) -> VortexResult<ArrayRef> {
79    let precision = array.decimal_dtype().precision();
80    let scale = array.decimal_dtype().scale();
81
82    let DataType::Decimal128(p, s) = data_type else {
83        vortex_bail!(
84            "Target Arrow type does not match Decimal source: {:?} ≠ {:?}",
85            data_type,
86            array.decimal_dtype()
87        );
88    };
89    if *p != precision || *s != scale {
90        vortex_bail!(
91            "Decimal128: precision {} and scale {} do not match expected ({}, {})",
92            precision,
93            scale,
94            p,
95            s
96        );
97    }
98
99    Ok(Arc::new(
100        Decimal128Array::new(
101            convert(array.buffer::<T>()).into_arrow_scalar_buffer(),
102            array.validity_mask()?.to_null_buffer(),
103        )
104        .with_precision_and_scale(precision, scale)?,
105    ))
106}
107
108fn convert_buffer<T: NativeDecimalType + AsPrimitive<i128>>(buffer: Buffer<T>) -> Buffer<i128> {
109    buffer.iter().map(|val| val.as_()).collect()
110}
111
112#[cfg(test)]
113mod tests {
114    use arrow_array::{Array, Decimal128Array};
115    use arrow_schema::DataType;
116    use vortex_buffer::buffer;
117    use vortex_dtype::DecimalDType;
118
119    use crate::arrays::DecimalArray;
120    use crate::compute::to_arrow;
121    use crate::validity::Validity;
122
123    #[test]
124    fn test_to_arrow() {
125        // Make a very simple i128 and i256 array.
126        let decimal_vortex = DecimalArray::new(
127            buffer![1i128, 2i128, 3i128, 4i128, 5i128],
128            DecimalDType::new(19, 2),
129            Validity::NonNullable,
130        );
131        let arrow = to_arrow(&decimal_vortex, &DataType::Decimal128(19, 2)).unwrap();
132        assert_eq!(arrow.data_type(), &DataType::Decimal128(19, 2));
133        let decimal_array = arrow.as_any().downcast_ref::<Decimal128Array>().unwrap();
134        assert_eq!(
135            decimal_array.values().as_ref(),
136            &[1i128, 2i128, 3i128, 4i128, 5i128]
137        );
138    }
139}
vortex_array/arrays/decimal/compute/to_arrow.rs

vortex_array/arrays/decimal/compute/
to_arrow.rs
