vortex_array/arrays/patched/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ExecutionCtx;
7use crate::array::ArrayView;
8use crate::array::OperationsVTable;
9use crate::arrays::PrimitiveArray;
10use crate::arrays::patched::Patched;
11use crate::arrays::patched::PatchedArrayExt;
12use crate::arrays::patched::PatchedArraySlotsExt;
13use crate::optimizer::ArrayOptimizer;
14use crate::scalar::Scalar;
15
16impl OperationsVTable<Patched> for Patched {
17    type ProbeState = ();
18
19    fn scalar_at(
20        array: ArrayView<'_, Patched>,
21        index: usize,
22        ctx: &mut ExecutionCtx,
23    ) -> VortexResult<Scalar> {
24        let chunk = (index + array.offset()) / 1024;
25
26        #[expect(
27            clippy::cast_possible_truncation,
28            reason = "N % 1024 always fits in u16"
29        )]
30        let chunk_index = ((index + array.offset()) % 1024) as u16;
31
32        let lane = (index + array.offset()) % array.n_lanes();
33
34        let range = array.lane_range(chunk, lane, ctx)?;
35
36        // Get the range of indices corresponding to the lane, potentially decoding them to avoid
37        // the overhead of repeated scalar_at calls.
38        let patch_indices = array
39            .patch_indices()
40            .slice(range.clone())?
41            .optimize()?
42            .execute::<PrimitiveArray>(ctx)?;
43
44        // NOTE: we do linear scan as lane has <= 32 patches, binary search would likely
45        //  be slower.
46        for (&patch_index, idx) in std::iter::zip(patch_indices.as_slice::<u16>(), range) {
47            if patch_index == chunk_index {
48                return array
49                    .patch_values()
50                    .execute_scalar(idx, ctx)?
51                    .cast(array.dtype());
52            }
53        }
54
55        // Otherwise, access the underlying value.
56        array.inner().execute_scalar(index, ctx)
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use vortex_buffer::buffer;
63    use vortex_session::VortexSession;
64
65    use crate::IntoArray;
66    use crate::VortexSessionExecute;
67    use crate::array_session;
68    use crate::arrays::Patched;
69    use crate::dtype::Nullability;
70    use crate::optimizer::ArrayOptimizer;
71    use crate::patches::Patches;
72    use crate::scalar::Scalar;
73
74    #[test]
75    fn test_simple() {
76        let mut ctx = array_session().create_execution_ctx();
77        let values = buffer![0u16; 1024].into_array();
78        let patches = Patches::new(
79            1024,
80            0,
81            buffer![1u32, 2, 3].into_array(),
82            buffer![1u16; 3].into_array(),
83            None,
84            &mut ctx,
85        )
86        .unwrap();
87
88        let session = VortexSession::empty();
89        let mut ctx = session.create_execution_ctx();
90
91        let array = Patched::from_array_and_patches(values, &patches, &mut ctx)
92            .unwrap()
93            .into_array();
94
95        assert_eq!(
96            array
97                .execute_scalar(0, &mut array_session().create_execution_ctx())
98                .unwrap(),
99            Scalar::primitive(0u16, Nullability::NonNullable)
100        );
101        assert_eq!(
102            array
103                .execute_scalar(1, &mut array_session().create_execution_ctx())
104                .unwrap(),
105            Scalar::primitive(1u16, Nullability::NonNullable)
106        );
107        assert_eq!(
108            array
109                .execute_scalar(2, &mut array_session().create_execution_ctx())
110                .unwrap(),
111            Scalar::primitive(1u16, Nullability::NonNullable)
112        );
113        assert_eq!(
114            array
115                .execute_scalar(3, &mut array_session().create_execution_ctx())
116                .unwrap(),
117            Scalar::primitive(1u16, Nullability::NonNullable)
118        );
119    }
120
121    #[test]
122    fn test_multi_chunk() {
123        let mut ctx = array_session().create_execution_ctx();
124        let values = buffer![0u16; 4096].into_array();
125        let patches = Patches::new(
126            4096,
127            0,
128            buffer![1u32, 2, 3].into_array(),
129            buffer![1u16; 3].into_array(),
130            None,
131            &mut ctx,
132        )
133        .unwrap();
134
135        let session = VortexSession::empty();
136        let mut ctx = session.create_execution_ctx();
137
138        let array = Patched::from_array_and_patches(values, &patches, &mut ctx)
139            .unwrap()
140            .into_array();
141
142        for index in 0..array.len() {
143            let value = array
144                .execute_scalar(index, &mut array_session().create_execution_ctx())
145                .unwrap();
146
147            if [1, 2, 3].contains(&index) {
148                assert_eq!(value, 1u16.into());
149            } else {
150                assert_eq!(value, 0u16.into());
151            }
152        }
153    }
154
155    #[test]
156    fn test_multi_chunk_sliced() {
157        let mut ctx = array_session().create_execution_ctx();
158        let values = buffer![0u16; 4096].into_array();
159        let patches = Patches::new(
160            4096,
161            0,
162            buffer![1u32, 2, 3].into_array(),
163            buffer![1u16; 3].into_array(),
164            None,
165            &mut ctx,
166        )
167        .unwrap();
168
169        let session = VortexSession::empty();
170        let mut ctx = session.create_execution_ctx();
171
172        let array = Patched::from_array_and_patches(values, &patches, &mut ctx)
173            .unwrap()
174            .into_array()
175            .slice(3..4096)
176            .unwrap()
177            .optimize()
178            .unwrap();
179
180        assert!(array.is::<Patched>());
181
182        assert_eq!(
183            array
184                .execute_scalar(0, &mut array_session().create_execution_ctx())
185                .unwrap(),
186            1u16.into()
187        );
188        for index in 1..array.len() {
189            assert_eq!(
190                array
191                    .execute_scalar(index, &mut array_session().create_execution_ctx())
192                    .unwrap(),
193                0u16.into()
194            );
195        }
196    }
197}
vortex_array/arrays/patched/vtable/operations.rs

vortex_array/arrays/patched/vtable/
operations.rs
