vortex_array/arrays/dict/compute/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod cast;
5mod compare;
6mod fill_null;
7pub(crate) mod is_constant;
8pub(crate) mod is_sorted;
9mod like;
10mod mask;
11pub(crate) mod min_max;
12pub(crate) mod rules;
13mod slice;
14
15use vortex_error::VortexResult;
16use vortex_mask::Mask;
17
18use super::Dict;
19use super::DictArray;
20use super::TakeExecute;
21use crate::ArrayRef;
22use crate::ExecutionCtx;
23use crate::IntoArray;
24use crate::array::ArrayView;
25use crate::arrays::dict::DictArraySlotsExt;
26use crate::arrays::filter::FilterReduce;
27
28impl TakeExecute for Dict {
29    fn take(
30        array: ArrayView<'_, Dict>,
31        indices: &ArrayRef,
32        _ctx: &mut ExecutionCtx,
33    ) -> VortexResult<Option<ArrayRef>> {
34        let codes = array.codes().take(indices.clone())?;
35
36        // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code
37        // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because
38        // selection can remove a value's last reference.
39        let taken = unsafe { DictArray::new_unchecked(codes, array.values().clone()) };
40
41        Ok(Some(taken.into_array()))
42    }
43}
44
45impl FilterReduce for Dict {
46    fn filter(array: ArrayView<'_, Dict>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
47        let codes = array.codes().filter(mask.clone())?;
48
49        // SAFETY: Selection preserves the integer code type and non-null code bounds, so every code
50        // indexes the unchanged values. `new_unchecked` resets `all_values_referenced` because
51        // selection can remove a value's last reference.
52        let filtered = unsafe { DictArray::new_unchecked(codes, array.values().clone()) };
53
54        Ok(Some(filtered.into_array()))
55    }
56}
57
58#[cfg(test)]
59mod test {
60    use std::sync::LazyLock;
61
62    use vortex_buffer::buffer;
63    use vortex_error::VortexResult;
64    use vortex_session::VortexSession;
65
66    use crate::ArrayRef;
67    use crate::IntoArray;
68    use crate::VortexSessionExecute;
69    use crate::arrays::ConstantArray;
70    use crate::arrays::PrimitiveArray;
71    use crate::arrays::VarBinArray;
72    use crate::arrays::VarBinViewArray;
73    use crate::assert_arrays_eq;
74    use crate::builders::dict::dict_encode;
75    use crate::builtins::ArrayBuiltins;
76    use crate::compute::conformance::filter::test_filter_conformance;
77    use crate::compute::conformance::mask::test_mask_conformance;
78    use crate::compute::conformance::take::test_take_conformance;
79    use crate::dtype::DType;
80    use crate::dtype::Nullability;
81    use crate::dtype::PType::I32;
82    use crate::scalar_fn::fns::operators::Operator;
83
84    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
85
86    #[test]
87    fn canonicalise_nullable_primitive() {
88        let mut ctx = SESSION.create_execution_ctx();
89        let values: Vec<Option<i32>> = (0..65)
90            .map(|i| match i % 3 {
91                0 => Some(42),
92                1 => Some(-9),
93                2 => None,
94                _ => unreachable!(),
95            })
96            .collect();
97
98        let dict = dict_encode(
99            &PrimitiveArray::from_option_iter(values.clone()).into_array(),
100            &mut SESSION.create_execution_ctx(),
101        )
102        .unwrap();
103
104        let expected = PrimitiveArray::from_option_iter(values);
105
106        assert_arrays_eq!(dict, expected, &mut ctx);
107    }
108
109    #[test]
110    fn canonicalise_non_nullable_primitive_32_unique_values() {
111        let mut ctx = SESSION.create_execution_ctx();
112        let unique_values: Vec<i32> = (0..32).collect();
113        let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 32]));
114
115        let dict = dict_encode(
116            &expected.clone().into_array(),
117            &mut SESSION.create_execution_ctx(),
118        )
119        .unwrap()
120        .into_array();
121
122        assert_arrays_eq!(dict, expected, &mut ctx);
123    }
124
125    #[test]
126    fn canonicalise_non_nullable_primitive_100_unique_values() {
127        let mut ctx = SESSION.create_execution_ctx();
128        let unique_values: Vec<i32> = (0..100).collect();
129        let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 100]));
130
131        let dict = dict_encode(
132            &expected.clone().into_array(),
133            &mut SESSION.create_execution_ctx(),
134        )
135        .unwrap()
136        .into_array();
137
138        assert_arrays_eq!(dict, expected, &mut ctx);
139    }
140
141    #[test]
142    fn canonicalise_nullable_varbin() -> VortexResult<()> {
143        let reference = VarBinViewArray::from_iter(
144            vec![Some("a"), Some("b"), None, Some("a"), None, Some("b")],
145            DType::Utf8(Nullability::Nullable),
146        );
147        assert_eq!(reference.len(), 6);
148        let mut ctx = SESSION.create_execution_ctx();
149        let dict = dict_encode(&reference.clone().into_array(), &mut ctx)?;
150        let flattened_dict = dict.into_array().execute::<VarBinViewArray>(&mut ctx)?;
151        let flattened_mask = flattened_dict
152            .validity()?
153            .execute_mask(flattened_dict.len(), &mut ctx)?;
154        let flattened_values = (0..flattened_dict.len())
155            .map(|i| {
156                flattened_mask
157                    .value(i)
158                    .then(|| flattened_dict.bytes_at(i).to_vec())
159            })
160            .collect::<Vec<_>>();
161        let reference_mask = reference
162            .validity()?
163            .execute_mask(reference.len(), &mut ctx)?;
164        let reference_values = (0..reference.len())
165            .map(|i| {
166                reference_mask
167                    .value(i)
168                    .then(|| reference.bytes_at(i).to_vec())
169            })
170            .collect::<Vec<_>>();
171        assert_eq!(flattened_values, reference_values);
172        Ok(())
173    }
174
175    fn sliced_dict_array() -> ArrayRef {
176        let reference = PrimitiveArray::from_option_iter([
177            Some(42),
178            Some(-9),
179            None,
180            Some(42),
181            Some(1),
182            Some(5),
183        ]);
184        let dict =
185            dict_encode(&reference.into_array(), &mut SESSION.create_execution_ctx()).unwrap();
186        dict.slice(1..4).unwrap()
187    }
188
189    #[test]
190    fn compare_sliced_dict() {
191        let mut ctx = SESSION.create_execution_ctx();
192        use crate::arrays::BoolArray;
193        let sliced = sliced_dict_array();
194        let compared = sliced
195            .binary(ConstantArray::new(42, 3).into_array(), Operator::Eq)
196            .unwrap();
197
198        let expected = BoolArray::from_iter([Some(false), None, Some(true)]);
199        assert_arrays_eq!(compared, expected.into_array(), &mut ctx);
200    }
201
202    #[test]
203    fn test_mask_dict_array() {
204        let array = dict_encode(
205            &buffer![2, 0, 2, 0, 10].into_array(),
206            &mut SESSION.create_execution_ctx(),
207        )
208        .unwrap();
209        test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
210
211        let array = dict_encode(
212            &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)])
213                .into_array(),
214            &mut SESSION.create_execution_ctx(),
215        )
216        .unwrap();
217        test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
218
219        let array = dict_encode(
220            &VarBinArray::from_iter(
221                [
222                    Some("hello"),
223                    None,
224                    Some("hello"),
225                    Some("good"),
226                    Some("good"),
227                ],
228                DType::Utf8(Nullability::Nullable),
229            )
230            .into_array(),
231            &mut SESSION.create_execution_ctx(),
232        )
233        .unwrap();
234        test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
235    }
236
237    #[test]
238    fn test_filter_dict_array() {
239        let array = dict_encode(
240            &buffer![2, 0, 2, 0, 10].into_array(),
241            &mut SESSION.create_execution_ctx(),
242        )
243        .unwrap();
244        test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
245
246        let array = dict_encode(
247            &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)])
248                .into_array(),
249            &mut SESSION.create_execution_ctx(),
250        )
251        .unwrap();
252        test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
253
254        let array = dict_encode(
255            &VarBinArray::from_iter(
256                [
257                    Some("hello"),
258                    None,
259                    Some("hello"),
260                    Some("good"),
261                    Some("good"),
262                ],
263                DType::Utf8(Nullability::Nullable),
264            )
265            .into_array(),
266            &mut SESSION.create_execution_ctx(),
267        )
268        .unwrap();
269        test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
270    }
271
272    #[test]
273    fn test_take_dict() {
274        let array = dict_encode(
275            &buffer![1, 2].into_array(),
276            &mut SESSION.create_execution_ctx(),
277        )
278        .unwrap();
279
280        let taken = array
281            .take(PrimitiveArray::from_option_iter([None::<i32>]).into_array())
282            .unwrap();
283        assert_eq!(taken.dtype(), &DType::Primitive(I32, Nullability::Nullable));
284        assert_arrays_eq!(
285            taken,
286            PrimitiveArray::from_option_iter([None::<i32>]),
287            &mut SESSION.create_execution_ctx()
288        );
289    }
290
291    #[test]
292    fn test_take_dict_conformance() {
293        let array = dict_encode(
294            &buffer![2, 0, 2, 0, 10].into_array(),
295            &mut SESSION.create_execution_ctx(),
296        )
297        .unwrap();
298        test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
299
300        let array = dict_encode(
301            &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)])
302                .into_array(),
303            &mut SESSION.create_execution_ctx(),
304        )
305        .unwrap();
306        test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
307
308        let array = dict_encode(
309            &VarBinArray::from_iter(
310                [
311                    Some("hello"),
312                    None,
313                    Some("hello"),
314                    Some("good"),
315                    Some("good"),
316                ],
317                DType::Utf8(Nullability::Nullable),
318            )
319            .into_array(),
320            &mut SESSION.create_execution_ctx(),
321        )
322        .unwrap();
323        test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use std::sync::LazyLock;
330
331    use rstest::rstest;
332    use vortex_buffer::buffer;
333    use vortex_session::VortexSession;
334
335    use crate::IntoArray;
336    use crate::VortexSessionExecute;
337    use crate::arrays::DictArray;
338    use crate::arrays::PrimitiveArray;
339    use crate::arrays::VarBinArray;
340    use crate::builders::dict::dict_encode;
341    use crate::compute::conformance::consistency::test_array_consistency;
342    use crate::dtype::DType;
343    use crate::dtype::Nullability;
344
345    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
346
347    #[rstest]
348    // Primitive arrays
349    #[case::dict_i32(dict_encode(&buffer![1i32, 2, 3, 2, 1].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
350    #[case::dict_nullable_codes(DictArray::try_new(
351        buffer![0u32, 1, 2, 2, 0].into_array(),
352        PrimitiveArray::from_option_iter([Some(10), Some(20), None]).into_array(),
353    ).unwrap())]
354    #[case::dict_nullable_values(dict_encode(
355        &PrimitiveArray::from_option_iter([Some(1i32), None, Some(2), Some(1), None]).into_array()
356    , &mut SESSION.create_execution_ctx()).unwrap())]
357    #[case::dict_u64(dict_encode(&buffer![100u64, 200, 100, 300, 200].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
358    // String arrays
359    #[case::dict_str(dict_encode(
360        &VarBinArray::from_iter(
361            ["hello", "world", "hello", "test", "world"].map(Some),
362            DType::Utf8(Nullability::NonNullable),
363        ).into_array()
364    , &mut SESSION.create_execution_ctx()).unwrap())]
365    #[case::dict_nullable_str(dict_encode(
366        &VarBinArray::from_iter(
367            [Some("hello"), None, Some("world"), Some("hello"), None],
368            DType::Utf8(Nullability::Nullable),
369        ).into_array()
370    , &mut SESSION.create_execution_ctx()).unwrap())]
371    // Edge cases
372    #[case::dict_single(dict_encode(&buffer![42i32].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
373    #[case::dict_all_same(dict_encode(&buffer![5i32, 5, 5, 5, 5].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
374    #[case::dict_large(dict_encode(&PrimitiveArray::from_iter((0..1000).map(|i| i % 10)).into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
375    fn test_dict_consistency(#[case] array: DictArray) {
376        test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx());
377    }
378}
vortex_array/arrays/dict/compute/mod.rs

vortex_array/arrays/dict/compute/
mod.rs
