vortex_array/arrays/decimal/compute/
rules.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_error::VortexResult;
7
8use crate::ArrayRef;
9use crate::IntoArray;
10use crate::array::ArrayView;
11use crate::arrays::Decimal;
12use crate::arrays::DecimalArray;
13use crate::arrays::Masked;
14use crate::arrays::slice::SliceReduce;
15use crate::arrays::slice::SliceReduceAdaptor;
16use crate::match_each_decimal_value_type;
17use crate::optimizer::rules::ArrayParentReduceRule;
18use crate::optimizer::rules::ParentRuleSet;
19use crate::scalar_fn::fns::cast::CastReduceAdaptor;
20use crate::scalar_fn::fns::mask::MaskReduceAdaptor;
21
22pub(crate) static RULES: ParentRuleSet<Decimal> = ParentRuleSet::new(&[
23    ParentRuleSet::lift(&DecimalMaskedValidityRule),
24    ParentRuleSet::lift(&CastReduceAdaptor(Decimal)),
25    ParentRuleSet::lift(&MaskReduceAdaptor(Decimal)),
26    ParentRuleSet::lift(&SliceReduceAdaptor(Decimal)),
27]);
28
29/// Rule to push down validity masking from MaskedArray parent into DecimalArray child.
30///
31/// When a DecimalArray is wrapped by a MaskedArray, this rule merges the mask's validity
32/// with the DecimalArray's existing validity, eliminating the need for the MaskedArray wrapper.
33#[derive(Default, Debug)]
34pub struct DecimalMaskedValidityRule;
35
36impl ArrayParentReduceRule<Decimal> for DecimalMaskedValidityRule {
37    type Parent = Masked;
38
39    fn reduce_parent(
40        &self,
41        array: ArrayView<'_, Decimal>,
42        parent: ArrayView<'_, Masked>,
43        _child_idx: usize,
44    ) -> VortexResult<Option<ArrayRef>> {
45        // Merge the parent's validity mask into the child's validity
46        // TODO(joe): make this lazy
47        let masked_array = match_each_decimal_value_type!(array.values_type(), |D| {
48            // SAFETY: Since we are only flipping some bits in the validity, all invariants that
49            // were upheld are still upheld.
50            unsafe {
51                DecimalArray::new_unchecked(
52                    array.buffer::<D>(),
53                    array.decimal_dtype(),
54                    array.validity()?.and(parent.validity()?)?,
55                )
56            }
57            .into_array()
58        });
59
60        Ok(Some(masked_array))
61    }
62}
63
64impl SliceReduce for Decimal {
65    fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
66        let byte_width = array.values_type().byte_width();
67        let byte_range = range.start * byte_width..range.end * byte_width;
68        let values = array.buffer_handle().slice(byte_range);
69        let validity = array.validity()?.slice(range)?;
70
71        // SAFETY: Slicing on element boundaries preserves the buffer alignment, values type,
72        // decimal precision and scale, and validity length invariants.
73        let result = unsafe {
74            DecimalArray::new_unchecked_handle(
75                values,
76                array.values_type(),
77                array.decimal_dtype(),
78                validity,
79            )
80            .into_array()
81        };
82        Ok(Some(result))
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use std::any::Any;
89    use std::ops::Range;
90    use std::sync::Arc;
91
92    use futures::future::BoxFuture;
93    use num_traits::AsPrimitive;
94    use rstest::rstest;
95    use vortex_buffer::Alignment;
96    use vortex_buffer::Buffer;
97    use vortex_buffer::ByteBuffer;
98    use vortex_error::VortexResult;
99    use vortex_error::vortex_bail;
100    use vortex_error::vortex_err;
101
102    use crate::IntoArray;
103    use crate::VortexSessionExecute;
104    use crate::array_session;
105    use crate::arrays::Decimal;
106    use crate::arrays::DecimalArray;
107    use crate::assert_arrays_eq;
108    use crate::buffer::BufferHandle;
109    use crate::buffer::DeviceBuffer;
110    use crate::dtype::DecimalDType;
111    use crate::dtype::DecimalType;
112    use crate::match_each_decimal_value_type;
113    use crate::validity::Validity;
114
115    // Host-backed storage exposes device slicing without allowing implicit host copies.
116    #[derive(Debug, PartialEq, Eq, Hash)]
117    struct TestDeviceBuffer(ByteBuffer);
118
119    impl DeviceBuffer for TestDeviceBuffer {
120        fn as_any(&self) -> &dyn Any {
121            self
122        }
123
124        fn len(&self) -> usize {
125            self.0.len()
126        }
127
128        fn alignment(&self) -> Alignment {
129            self.0.alignment()
130        }
131
132        fn copy_to_host_sync(&self, _alignment: Alignment) -> VortexResult<ByteBuffer> {
133            vortex_bail!("decimal slicing must not copy device values to the host")
134        }
135
136        fn copy_to_host(
137            &self,
138            _alignment: Alignment,
139        ) -> VortexResult<BoxFuture<'static, VortexResult<ByteBuffer>>> {
140            vortex_bail!("decimal slicing must not copy device values to the host")
141        }
142
143        fn slice(&self, range: Range<usize>) -> Arc<dyn DeviceBuffer> {
144            Arc::new(Self(self.0.slice(range)))
145        }
146
147        fn aligned(self: Arc<Self>, alignment: Alignment) -> VortexResult<Arc<dyn DeviceBuffer>> {
148            assert!(self.alignment().is_aligned_to(alignment));
149            Ok(self)
150        }
151    }
152
153    // Exercise composed byte offsets at every storage width without allowing a host copy.
154    #[rstest]
155    fn test_slice_buffer_handle(
156        #[values(
157            DecimalType::I8,
158            DecimalType::I16,
159            DecimalType::I32,
160            DecimalType::I64,
161            DecimalType::I128,
162            DecimalType::I256
163        )]
164        values_type: DecimalType,
165    ) -> VortexResult<()> {
166        let mut ctx = array_session().create_execution_ctx();
167        let bytes = match_each_decimal_value_type!(values_type, |D| {
168            Buffer::<D>::from_iter([-12i8, 34, -56, 78, 90, 12].map(|value| value.as_()))
169                .into_byte_buffer()
170        });
171        let handle = BufferHandle::new_device(Arc::new(TestDeviceBuffer(bytes.clone())));
172        let validity = Validity::from_iter([true, false, true, false, true, true]);
173        let array = DecimalArray::try_new_handle(
174            handle,
175            values_type,
176            DecimalDType::new(3, 1),
177            validity.clone(),
178        )?
179        .into_array();
180        let sliced = array.slice(1..5)?;
181
182        let range = 1..3;
183        let nested = sliced.slice(range.clone())?;
184        let decimal = nested.as_::<Decimal>();
185        let byte_width = values_type.byte_width();
186        let original_range = range.start + 1..range.end + 1;
187        let expected_bytes =
188            bytes.slice(original_range.start * byte_width..original_range.end * byte_width);
189        let handle = decimal.buffer_handle();
190
191        assert_eq!(nested.len(), range.len());
192        assert_eq!(nested.dtype(), array.dtype());
193        assert_eq!(decimal.values_type(), values_type);
194        assert_eq!(handle.len(), range.len() * byte_width);
195        assert!(handle.is_on_device());
196        assert!(handle.is_aligned_to(bytes.alignment()));
197
198        let actual_bytes = &handle
199            .as_device()
200            .as_any()
201            .downcast_ref::<TestDeviceBuffer>()
202            .ok_or_else(|| vortex_err!("expected TestDeviceBuffer"))?
203            .0;
204        assert_eq!(actual_bytes, &expected_bytes);
205        assert_eq!(actual_bytes.as_ptr(), expected_bytes.as_ptr());
206        assert_arrays_eq!(
207            decimal.validity()?.to_array(nested.len()),
208            validity.to_array(array.len()).slice(original_range)?,
209            &mut ctx
210        );
211
212        Ok(())
213    }
214}
vortex_array/arrays/decimal/compute/rules.rs

vortex_array/arrays/decimal/compute/
rules.rs
