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        let actual = dict
104            .into_array()
105            .execute::<PrimitiveArray>(&mut SESSION.create_execution_ctx())
106            .unwrap();
107
108        let expected = PrimitiveArray::from_option_iter(values);
109
110        assert_arrays_eq!(actual, expected, &mut ctx);
111    }
112
113    #[test]
114    fn canonicalise_non_nullable_primitive_32_unique_values() {
115        let mut ctx = SESSION.create_execution_ctx();
116        let unique_values: Vec<i32> = (0..32).collect();
117        let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 32]));
118
119        let dict = dict_encode(
120            &expected.clone().into_array(),
121            &mut SESSION.create_execution_ctx(),
122        )
123        .unwrap()
124        .into_array();
125
126        let actual = dict
127            .execute::<PrimitiveArray>(&mut SESSION.create_execution_ctx())
128            .unwrap();
129
130        assert_arrays_eq!(actual, expected, &mut ctx);
131    }
132
133    #[test]
134    fn canonicalise_non_nullable_primitive_100_unique_values() {
135        let mut ctx = SESSION.create_execution_ctx();
136        let unique_values: Vec<i32> = (0..100).collect();
137        let expected = PrimitiveArray::from_iter((0..1000).map(|i| unique_values[i % 100]));
138
139        let dict = dict_encode(
140            &expected.clone().into_array(),
141            &mut SESSION.create_execution_ctx(),
142        )
143        .unwrap()
144        .into_array();
145
146        let actual = dict
147            .execute::<PrimitiveArray>(&mut SESSION.create_execution_ctx())
148            .unwrap();
149
150        assert_arrays_eq!(actual, expected, &mut ctx);
151    }
152
153    #[test]
154    fn canonicalise_nullable_varbin() -> VortexResult<()> {
155        let reference = VarBinViewArray::from_iter(
156            vec![Some("a"), Some("b"), None, Some("a"), None, Some("b")],
157            DType::Utf8(Nullability::Nullable),
158        );
159        assert_eq!(reference.len(), 6);
160        let mut ctx = SESSION.create_execution_ctx();
161        let dict = dict_encode(&reference.clone().into_array(), &mut ctx)?;
162        let flattened_dict = dict.into_array().execute::<VarBinViewArray>(&mut ctx)?;
163        let flattened_mask = flattened_dict
164            .validity()?
165            .execute_mask(flattened_dict.len(), &mut ctx)?;
166        let flattened_values = (0..flattened_dict.len())
167            .map(|i| {
168                flattened_mask
169                    .value(i)
170                    .then(|| flattened_dict.bytes_at(i).to_vec())
171            })
172            .collect::<Vec<_>>();
173        let reference_mask = reference
174            .validity()?
175            .execute_mask(reference.len(), &mut ctx)?;
176        let reference_values = (0..reference.len())
177            .map(|i| {
178                reference_mask
179                    .value(i)
180                    .then(|| reference.bytes_at(i).to_vec())
181            })
182            .collect::<Vec<_>>();
183        assert_eq!(flattened_values, reference_values);
184        Ok(())
185    }
186
187    fn sliced_dict_array() -> ArrayRef {
188        let reference = PrimitiveArray::from_option_iter([
189            Some(42),
190            Some(-9),
191            None,
192            Some(42),
193            Some(1),
194            Some(5),
195        ]);
196        let dict =
197            dict_encode(&reference.into_array(), &mut SESSION.create_execution_ctx()).unwrap();
198        dict.slice(1..4).unwrap()
199    }
200
201    #[test]
202    fn compare_sliced_dict() {
203        let mut ctx = SESSION.create_execution_ctx();
204        use crate::arrays::BoolArray;
205        let sliced = sliced_dict_array();
206        let compared = sliced
207            .binary(ConstantArray::new(42, 3).into_array(), Operator::Eq)
208            .unwrap();
209
210        let expected = BoolArray::from_iter([Some(false), None, Some(true)]);
211        assert_arrays_eq!(compared, expected.into_array(), &mut ctx);
212    }
213
214    #[test]
215    fn test_mask_dict_array() {
216        let array = dict_encode(
217            &buffer![2, 0, 2, 0, 10].into_array(),
218            &mut SESSION.create_execution_ctx(),
219        )
220        .unwrap();
221        test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
222
223        let array = dict_encode(
224            &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)])
225                .into_array(),
226            &mut SESSION.create_execution_ctx(),
227        )
228        .unwrap();
229        test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
230
231        let array = dict_encode(
232            &VarBinArray::from_iter(
233                [
234                    Some("hello"),
235                    None,
236                    Some("hello"),
237                    Some("good"),
238                    Some("good"),
239                ],
240                DType::Utf8(Nullability::Nullable),
241            )
242            .into_array(),
243            &mut SESSION.create_execution_ctx(),
244        )
245        .unwrap();
246        test_mask_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
247    }
248
249    #[test]
250    fn test_filter_dict_array() {
251        let array = dict_encode(
252            &buffer![2, 0, 2, 0, 10].into_array(),
253            &mut SESSION.create_execution_ctx(),
254        )
255        .unwrap();
256        test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
257
258        let array = dict_encode(
259            &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)])
260                .into_array(),
261            &mut SESSION.create_execution_ctx(),
262        )
263        .unwrap();
264        test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
265
266        let array = dict_encode(
267            &VarBinArray::from_iter(
268                [
269                    Some("hello"),
270                    None,
271                    Some("hello"),
272                    Some("good"),
273                    Some("good"),
274                ],
275                DType::Utf8(Nullability::Nullable),
276            )
277            .into_array(),
278            &mut SESSION.create_execution_ctx(),
279        )
280        .unwrap();
281        test_filter_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
282    }
283
284    #[test]
285    fn test_take_dict() {
286        let array = dict_encode(
287            &buffer![1, 2].into_array(),
288            &mut SESSION.create_execution_ctx(),
289        )
290        .unwrap();
291
292        assert_eq!(
293            array
294                .take(PrimitiveArray::from_option_iter([Option::<i32>::None]).into_array())
295                .unwrap()
296                .dtype(),
297            &DType::Primitive(I32, Nullability::Nullable)
298        );
299    }
300
301    #[test]
302    fn test_take_dict_conformance() {
303        let array = dict_encode(
304            &buffer![2, 0, 2, 0, 10].into_array(),
305            &mut SESSION.create_execution_ctx(),
306        )
307        .unwrap();
308        test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
309
310        let array = dict_encode(
311            &PrimitiveArray::from_option_iter([Some(2), None, Some(2), Some(0), Some(10)])
312                .into_array(),
313            &mut SESSION.create_execution_ctx(),
314        )
315        .unwrap();
316        test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
317
318        let array = dict_encode(
319            &VarBinArray::from_iter(
320                [
321                    Some("hello"),
322                    None,
323                    Some("hello"),
324                    Some("good"),
325                    Some("good"),
326                ],
327                DType::Utf8(Nullability::Nullable),
328            )
329            .into_array(),
330            &mut SESSION.create_execution_ctx(),
331        )
332        .unwrap();
333        test_take_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
334    }
335}
336
337#[cfg(test)]
338mod tests {
339    use std::sync::LazyLock;
340
341    use rstest::rstest;
342    use vortex_buffer::buffer;
343    use vortex_session::VortexSession;
344
345    use crate::IntoArray;
346    use crate::VortexSessionExecute;
347    use crate::arrays::DictArray;
348    use crate::arrays::PrimitiveArray;
349    use crate::arrays::VarBinArray;
350    use crate::builders::dict::dict_encode;
351    use crate::compute::conformance::consistency::test_array_consistency;
352    use crate::dtype::DType;
353    use crate::dtype::Nullability;
354
355    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
356
357    #[rstest]
358    // Primitive arrays
359    #[case::dict_i32(dict_encode(&buffer![1i32, 2, 3, 2, 1].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
360    #[case::dict_nullable_codes(DictArray::try_new(
361        buffer![0u32, 1, 2, 2, 0].into_array(),
362        PrimitiveArray::from_option_iter([Some(10), Some(20), None]).into_array(),
363    ).unwrap())]
364    #[case::dict_nullable_values(dict_encode(
365        &PrimitiveArray::from_option_iter([Some(1i32), None, Some(2), Some(1), None]).into_array()
366    , &mut SESSION.create_execution_ctx()).unwrap())]
367    #[case::dict_u64(dict_encode(&buffer![100u64, 200, 100, 300, 200].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
368    // String arrays
369    #[case::dict_str(dict_encode(
370        &VarBinArray::from_iter(
371            ["hello", "world", "hello", "test", "world"].map(Some),
372            DType::Utf8(Nullability::NonNullable),
373        ).into_array()
374    , &mut SESSION.create_execution_ctx()).unwrap())]
375    #[case::dict_nullable_str(dict_encode(
376        &VarBinArray::from_iter(
377            [Some("hello"), None, Some("world"), Some("hello"), None],
378            DType::Utf8(Nullability::Nullable),
379        ).into_array()
380    , &mut SESSION.create_execution_ctx()).unwrap())]
381    // Edge cases
382    #[case::dict_single(dict_encode(&buffer![42i32].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
383    #[case::dict_all_same(dict_encode(&buffer![5i32, 5, 5, 5, 5].into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
384    #[case::dict_large(dict_encode(&PrimitiveArray::from_iter((0..1000).map(|i| i % 10)).into_array(), &mut SESSION.create_execution_ctx()).unwrap())]
385    fn test_dict_consistency(#[case] array: DictArray) {
386        test_array_consistency(&array.into_array(), &mut SESSION.create_execution_ctx());
387    }
388}
vortex_array/arrays/dict/compute/mod.rs

vortex_array/arrays/dict/compute/
mod.rs
