vortex_array/arrays/decimal/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_buffer::Alignment;
5use vortex_buffer::Buffer;
6use vortex_dtype::DType;
7use vortex_dtype::NativeDecimalType;
8use vortex_dtype::PrecisionScale;
9use vortex_dtype::match_each_decimal_value_type;
10use vortex_error::VortexExpect;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_ensure;
14use vortex_scalar::DecimalType;
15use vortex_vector::decimal::DVector;
16
17use crate::ArrayRef;
18use crate::DeserializeMetadata;
19use crate::ProstMetadata;
20use crate::SerializeMetadata;
21use crate::arrays::DecimalArray;
22use crate::buffer::BufferHandle;
23use crate::executor::ExecutionCtx;
24use crate::serde::ArrayChildren;
25use crate::validity::Validity;
26use crate::vtable;
27use crate::vtable::ArrayVTableExt;
28use crate::vtable::NotSupported;
29use crate::vtable::VTable;
30use crate::vtable::ValidityVTableFromValidityHelper;
31
32mod array;
33mod canonical;
34mod operations;
35pub mod rules;
36mod validity;
37mod visitor;
38
39pub use rules::DecimalMaskedValidityRule;
40use vortex_vector::Vector;
41
42use crate::arrays::decimal::vtable::rules::RULES;
43use crate::vtable::ArrayId;
44use crate::vtable::ArrayVTable;
45
46vtable!(Decimal);
47
48// The type of the values can be determined by looking at the type info...right?
49#[derive(prost::Message)]
50pub struct DecimalMetadata {
51    #[prost(enumeration = "DecimalType", tag = "1")]
52    pub(super) values_type: i32,
53}
54
55impl VTable for DecimalVTable {
56    type Array = DecimalArray;
57
58    type Metadata = ProstMetadata<DecimalMetadata>;
59
60    type ArrayVTable = Self;
61    type CanonicalVTable = Self;
62    type OperationsVTable = Self;
63    type ValidityVTable = ValidityVTableFromValidityHelper;
64    type VisitorVTable = Self;
65    type ComputeVTable = NotSupported;
66    type EncodeVTable = NotSupported;
67
68    fn id(&self) -> ArrayId {
69        ArrayId::new_ref("vortex.decimal")
70    }
71
72    fn encoding(_array: &Self::Array) -> ArrayVTable {
73        DecimalVTable.as_vtable()
74    }
75
76    fn metadata(array: &DecimalArray) -> VortexResult<Self::Metadata> {
77        Ok(ProstMetadata(DecimalMetadata {
78            values_type: array.values_type() as i32,
79        }))
80    }
81
82    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
83        Ok(Some(metadata.serialize()))
84    }
85
86    fn deserialize(bytes: &[u8]) -> VortexResult<Self::Metadata> {
87        let metadata = ProstMetadata::<DecimalMetadata>::deserialize(bytes)?;
88        Ok(ProstMetadata(metadata))
89    }
90
91    fn build(
92        &self,
93        dtype: &DType,
94        len: usize,
95        metadata: &Self::Metadata,
96        buffers: &[BufferHandle],
97        children: &dyn ArrayChildren,
98    ) -> VortexResult<DecimalArray> {
99        if buffers.len() != 1 {
100            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
101        }
102        let buffer = buffers[0].clone().try_to_bytes()?;
103
104        let validity = if children.is_empty() {
105            Validity::from(dtype.nullability())
106        } else if children.len() == 1 {
107            let validity = children.get(0, &Validity::DTYPE, len)?;
108            Validity::Array(validity)
109        } else {
110            vortex_bail!("Expected 0 or 1 child, got {}", children.len());
111        };
112
113        let Some(decimal_dtype) = dtype.as_decimal_opt() else {
114            vortex_bail!("Expected Decimal dtype, got {:?}", dtype)
115        };
116
117        match_each_decimal_value_type!(metadata.values_type(), |D| {
118            // Check and reinterpret-cast the buffer
119            vortex_ensure!(
120                buffer.is_aligned(Alignment::of::<D>()),
121                "DecimalArray buffer not aligned for values type {:?}",
122                D::DECIMAL_TYPE
123            );
124            let buffer = Buffer::<D>::from_byte_buffer(buffer);
125            DecimalArray::try_new::<D>(buffer, *decimal_dtype, validity)
126        })
127    }
128
129    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
130        vortex_ensure!(
131            children.len() <= 1,
132            "DecimalArray expects 0 or 1 child (validity), got {}",
133            children.len()
134        );
135
136        if children.is_empty() {
137            array.validity = Validity::from(array.dtype.nullability());
138        } else {
139            array.validity = Validity::Array(
140                children
141                    .into_iter()
142                    .next()
143                    .vortex_expect("children length already validated"),
144            );
145        }
146        Ok(())
147    }
148
149    fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<Vector> {
150        use vortex_dtype::BigCast;
151
152        match_each_decimal_value_type!(array.values_type(), |D| {
153            // TODO(ngates): we probably shouldn't convert here... Do we allow larger P/S for a
154            //  given physical type, because we know that our values actually fit?
155            let min_value_type = DecimalType::smallest_decimal_value_type(&array.decimal_dtype());
156            match_each_decimal_value_type!(min_value_type, |E| {
157                let decimal_dtype = array.decimal_dtype();
158                let buffer = array.buffer::<D>();
159                let validity_mask = array.validity_mask();
160
161                // Copy from D to E, possibly widening, possibly narrowing
162                let values = Buffer::<E>::from_trusted_len_iter(
163                    buffer
164                        .iter()
165                        .map(|d| <E as BigCast>::from(*d).vortex_expect("Decimal cast failed")),
166                );
167
168                Ok(unsafe {
169                    DVector::<E>::new_unchecked(
170                        // TODO(ngates): this is too small?
171                        PrecisionScale::new_unchecked(
172                            decimal_dtype.precision(),
173                            decimal_dtype.scale(),
174                        ),
175                        values,
176                        validity_mask,
177                    )
178                }
179                .into())
180            })
181        })
182    }
183
184    fn reduce_parent(
185        array: &Self::Array,
186        parent: &ArrayRef,
187        child_idx: usize,
188    ) -> VortexResult<Option<ArrayRef>> {
189        RULES.evaluate(array, parent, child_idx)
190    }
191}
192
193#[derive(Debug)]
194pub struct DecimalVTable;
195
196#[cfg(test)]
197mod tests {
198    use vortex_buffer::ByteBufferMut;
199    use vortex_buffer::buffer;
200    use vortex_dtype::DecimalDType;
201
202    use crate::ArrayContext;
203    use crate::IntoArray;
204    use crate::arrays::DecimalArray;
205    use crate::arrays::DecimalVTable;
206    use crate::serde::ArrayParts;
207    use crate::serde::SerializeOptions;
208    use crate::validity::Validity;
209    use crate::vtable::ArrayVTableExt;
210
211    #[test]
212    fn test_array_serde() {
213        let array = DecimalArray::new(
214            buffer![100i128, 200i128, 300i128, 400i128, 500i128],
215            DecimalDType::new(10, 2),
216            Validity::NonNullable,
217        );
218        let dtype = array.dtype().clone();
219        let ctx = ArrayContext::empty().with(DecimalVTable.as_vtable());
220        let out = array
221            .into_array()
222            .serialize(&ctx, &SerializeOptions::default())
223            .unwrap();
224        // Concat into a single buffer
225        let mut concat = ByteBufferMut::empty();
226        for buf in out {
227            concat.extend_from_slice(buf.as_ref());
228        }
229
230        let concat = concat.freeze();
231
232        let parts = ArrayParts::try_from(concat).unwrap();
233
234        let decoded = parts.decode(&ctx, &dtype, 5).unwrap();
235        assert!(decoded.is::<DecimalVTable>());
236    }
237}
vortex_array/arrays/decimal/vtable/mod.rs

vortex_array/arrays/decimal/vtable/
mod.rs
