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)?;
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 values = buffer![0u16; 1024].into_array();
77        let patches = Patches::new(
78            1024,
79            0,
80            buffer![1u32, 2, 3].into_array(),
81            buffer![1u16; 3].into_array(),
82            None,
83        )
84        .unwrap();
85
86        let session = VortexSession::empty();
87        let mut ctx = session.create_execution_ctx();
88
89        let array = Patched::from_array_and_patches(values, &patches, &mut ctx)
90            .unwrap()
91            .into_array();
92
93        assert_eq!(
94            array
95                .execute_scalar(0, &mut array_session().create_execution_ctx())
96                .unwrap(),
97            Scalar::primitive(0u16, Nullability::NonNullable)
98        );
99        assert_eq!(
100            array
101                .execute_scalar(1, &mut array_session().create_execution_ctx())
102                .unwrap(),
103            Scalar::primitive(1u16, Nullability::NonNullable)
104        );
105        assert_eq!(
106            array
107                .execute_scalar(2, &mut array_session().create_execution_ctx())
108                .unwrap(),
109            Scalar::primitive(1u16, Nullability::NonNullable)
110        );
111        assert_eq!(
112            array
113                .execute_scalar(3, &mut array_session().create_execution_ctx())
114                .unwrap(),
115            Scalar::primitive(1u16, Nullability::NonNullable)
116        );
117    }
118
119    #[test]
120    fn test_multi_chunk() {
121        let values = buffer![0u16; 4096].into_array();
122        let patches = Patches::new(
123            4096,
124            0,
125            buffer![1u32, 2, 3].into_array(),
126            buffer![1u16; 3].into_array(),
127            None,
128        )
129        .unwrap();
130
131        let session = VortexSession::empty();
132        let mut ctx = session.create_execution_ctx();
133
134        let array = Patched::from_array_and_patches(values, &patches, &mut ctx)
135            .unwrap()
136            .into_array();
137
138        for index in 0..array.len() {
139            let value = array
140                .execute_scalar(index, &mut array_session().create_execution_ctx())
141                .unwrap();
142
143            if [1, 2, 3].contains(&index) {
144                assert_eq!(value, 1u16.into());
145            } else {
146                assert_eq!(value, 0u16.into());
147            }
148        }
149    }
150
151    #[test]
152    fn test_multi_chunk_sliced() {
153        let values = buffer![0u16; 4096].into_array();
154        let patches = Patches::new(
155            4096,
156            0,
157            buffer![1u32, 2, 3].into_array(),
158            buffer![1u16; 3].into_array(),
159            None,
160        )
161        .unwrap();
162
163        let session = VortexSession::empty();
164        let mut ctx = session.create_execution_ctx();
165
166        let array = Patched::from_array_and_patches(values, &patches, &mut ctx)
167            .unwrap()
168            .into_array()
169            .slice(3..4096)
170            .unwrap()
171            .optimize()
172            .unwrap();
173
174        assert!(array.is::<Patched>());
175
176        assert_eq!(
177            array
178                .execute_scalar(0, &mut array_session().create_execution_ctx())
179                .unwrap(),
180            1u16.into()
181        );
182        for index in 1..array.len() {
183            assert_eq!(
184                array
185                    .execute_scalar(index, &mut array_session().create_execution_ctx())
186                    .unwrap(),
187                0u16.into()
188            );
189        }
190    }
191}
vortex_array/arrays/patched/vtable/operations.rs

vortex_array/arrays/patched/vtable/
operations.rs
