vortex_dict/
array.rs

1use std::fmt::Debug;
2
3use arrow_buffer::BooleanBuffer;
4use vortex_array::compute::{cast, take};
5use vortex_array::stats::{ArrayStats, StatsSetRef};
6use vortex_array::vtable::{ArrayVTable, CanonicalVTable, NotSupported, VTable, ValidityVTable};
7use vortex_array::{
8    Array, ArrayRef, Canonical, EncodingId, EncodingRef, IntoArray, ToCanonical, vtable,
9};
10use vortex_dtype::{DType, match_each_integer_ptype};
11use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
12use vortex_mask::{AllOr, Mask};
13
14vtable!(Dict);
15
16impl VTable for DictVTable {
17    type Array = DictArray;
18    type Encoding = DictEncoding;
19
20    type ArrayVTable = Self;
21    type CanonicalVTable = Self;
22    type OperationsVTable = Self;
23    type ValidityVTable = Self;
24    type VisitorVTable = Self;
25    type ComputeVTable = NotSupported;
26    type EncodeVTable = Self;
27    type SerdeVTable = Self;
28
29    fn id(_encoding: &Self::Encoding) -> EncodingId {
30        EncodingId::new_ref("vortex.dict")
31    }
32
33    fn encoding(_array: &Self::Array) -> EncodingRef {
34        EncodingRef::new_ref(DictEncoding.as_ref())
35    }
36}
37
38#[derive(Debug, Clone)]
39pub struct DictArray {
40    codes: ArrayRef,
41    values: ArrayRef,
42    stats_set: ArrayStats,
43}
44
45#[derive(Clone, Debug)]
46pub struct DictEncoding;
47
48impl DictArray {
49    pub fn try_new(mut codes: ArrayRef, values: ArrayRef) -> VortexResult<Self> {
50        if !codes.dtype().is_unsigned_int() {
51            vortex_bail!(MismatchedTypes: "unsigned int", codes.dtype());
52        }
53
54        let dtype = values.dtype();
55        if dtype.is_nullable() {
56            // If the values are nullable, we force codes to be nullable as well.
57            codes = cast(&codes, &codes.dtype().as_nullable())?;
58        } else {
59            // If the values are non-nullable, we assert the codes are non-nullable as well.
60            if codes.dtype().is_nullable() {
61                vortex_bail!("Cannot have nullable codes for non-nullable dict array");
62            }
63        }
64        assert_eq!(
65            codes.dtype().nullability(),
66            values.dtype().nullability(),
67            "Mismatched nullability between codes and values"
68        );
69
70        Ok(Self {
71            codes,
72            values,
73            stats_set: Default::default(),
74        })
75    }
76
77    #[inline]
78    pub fn codes(&self) -> &ArrayRef {
79        &self.codes
80    }
81
82    #[inline]
83    pub fn values(&self) -> &ArrayRef {
84        &self.values
85    }
86}
87
88impl ArrayVTable<DictVTable> for DictVTable {
89    fn len(array: &DictArray) -> usize {
90        array.codes.len()
91    }
92
93    fn dtype(array: &DictArray) -> &DType {
94        array.values.dtype()
95    }
96
97    fn stats(array: &DictArray) -> StatsSetRef<'_> {
98        array.stats_set.to_ref(array.as_ref())
99    }
100}
101
102impl CanonicalVTable<DictVTable> for DictVTable {
103    fn canonicalize(array: &DictArray) -> VortexResult<Canonical> {
104        match array.dtype() {
105            // NOTE: Utf8 and Binary will decompress into VarBinViewArray, which requires a full
106            // decompression to construct the views child array.
107            // For this case, it is *always* faster to decompress the values first and then create
108            // copies of the view pointers.
109            DType::Utf8(_) | DType::Binary(_) => {
110                let canonical_values: ArrayRef = array.values().to_canonical()?.into_array();
111                take(&canonical_values, array.codes())?.to_canonical()
112            }
113            _ => take(array.values(), array.codes())?.to_canonical(),
114        }
115    }
116}
117
118impl ValidityVTable<DictVTable> for DictVTable {
119    fn is_valid(array: &DictArray, index: usize) -> VortexResult<bool> {
120        let scalar = array.codes().scalar_at(index).map_err(|err| {
121            err.with_context(format!("Failed to get index {index} from DictArray codes"))
122        })?;
123
124        if scalar.is_null() {
125            return Ok(false);
126        };
127        let values_index: usize = scalar
128            .as_ref()
129            .try_into()
130            .vortex_expect("Failed to convert dictionary code to usize");
131        array.values().is_valid(values_index)
132    }
133
134    fn all_valid(array: &DictArray) -> VortexResult<bool> {
135        Ok(array.codes().all_valid()? && array.values().all_valid()?)
136    }
137
138    fn all_invalid(array: &DictArray) -> VortexResult<bool> {
139        Ok(array.codes().all_invalid()? || array.values().all_invalid()?)
140    }
141
142    fn validity_mask(array: &DictArray) -> VortexResult<Mask> {
143        let codes_validity = array.codes().validity_mask()?;
144        match codes_validity.boolean_buffer() {
145            AllOr::All => {
146                let primitive_codes = array.codes().to_primitive()?;
147                let values_mask = array.values().validity_mask()?;
148                let is_valid_buffer = match_each_integer_ptype!(primitive_codes.ptype(), |$P| {
149                    let codes_slice = primitive_codes.as_slice::<$P>();
150                    BooleanBuffer::collect_bool(array.len(), |idx| {
151                       values_mask.value(codes_slice[idx] as usize)
152                    })
153                });
154                Ok(Mask::from_buffer(is_valid_buffer))
155            }
156            AllOr::None => Ok(Mask::AllFalse(array.len())),
157            AllOr::Some(validity_buff) => {
158                let primitive_codes = array.codes().to_primitive()?;
159                let values_mask = array.values().validity_mask()?;
160                let is_valid_buffer = match_each_integer_ptype!(primitive_codes.ptype(), |$P| {
161                    let codes_slice = primitive_codes.as_slice::<$P>();
162                    BooleanBuffer::collect_bool(array.len(), |idx| {
163                       validity_buff.value(idx) && values_mask.value(codes_slice[idx] as usize)
164                    })
165                });
166                Ok(Mask::from_buffer(is_valid_buffer))
167            }
168        }
169    }
170}
171
172#[cfg(test)]
173mod test {
174    use arrow_buffer::BooleanBuffer;
175    use rand::distr::{Distribution, StandardUniform};
176    use rand::prelude::StdRng;
177    use rand::{Rng, SeedableRng};
178    use vortex_array::arrays::{ChunkedArray, PrimitiveArray};
179    use vortex_array::builders::builder_with_capacity;
180    use vortex_array::validity::Validity;
181    use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
182    use vortex_buffer::buffer;
183    use vortex_dtype::Nullability::NonNullable;
184    use vortex_dtype::{DType, NativePType, PType};
185    use vortex_error::{VortexExpect, VortexUnwrap, vortex_panic};
186    use vortex_mask::AllOr;
187
188    use crate::DictArray;
189
190    #[test]
191    fn nullable_codes_validity() {
192        let dict = DictArray::try_new(
193            PrimitiveArray::new(
194                buffer![0u32, 1, 2, 2, 1],
195                Validity::from(BooleanBuffer::from(vec![true, false, true, false, true])),
196            )
197            .into_array(),
198            PrimitiveArray::new(buffer![3, 6, 9], Validity::AllValid).into_array(),
199        )
200        .unwrap();
201        let mask = dict.validity_mask().unwrap();
202        let AllOr::Some(indices) = mask.indices() else {
203            vortex_panic!("Expected indices from mask")
204        };
205        assert_eq!(indices, [0, 2, 4]);
206    }
207
208    #[test]
209    fn nullable_values_validity() {
210        let dict = DictArray::try_new(
211            buffer![0u32, 1, 2, 2, 1].into_array(),
212            PrimitiveArray::new(
213                buffer![3, 6, 9],
214                Validity::from(BooleanBuffer::from(vec![true, false, false])),
215            )
216            .into_array(),
217        )
218        .unwrap();
219        let mask = dict.validity_mask().unwrap();
220        let AllOr::Some(indices) = mask.indices() else {
221            vortex_panic!("Expected indices from mask")
222        };
223        assert_eq!(indices, [0]);
224    }
225
226    #[test]
227    fn nullable_codes_and_values() {
228        let dict = DictArray::try_new(
229            PrimitiveArray::new(
230                buffer![0u32, 1, 2, 2, 1],
231                Validity::from(BooleanBuffer::from(vec![true, false, true, false, true])),
232            )
233            .into_array(),
234            PrimitiveArray::new(
235                buffer![3, 6, 9],
236                Validity::from(BooleanBuffer::from(vec![false, true, true])),
237            )
238            .into_array(),
239        )
240        .unwrap();
241        let mask = dict.validity_mask().unwrap();
242        let AllOr::Some(indices) = mask.indices() else {
243            vortex_panic!("Expected indices from mask")
244        };
245        assert_eq!(indices, [2, 4]);
246    }
247
248    fn make_dict_primitive_chunks<T: NativePType, U: NativePType>(
249        len: usize,
250        unique_values: usize,
251        chunk_count: usize,
252    ) -> ArrayRef
253    where
254        StandardUniform: Distribution<T>,
255    {
256        let mut rng = StdRng::seed_from_u64(0);
257
258        (0..chunk_count)
259            .map(|_| {
260                let values = (0..unique_values)
261                    .map(|_| rng.random::<T>())
262                    .collect::<PrimitiveArray>();
263                let codes = (0..len)
264                    .map(|_| {
265                        U::from(rng.random_range(0..unique_values)).vortex_expect("valid value")
266                    })
267                    .collect::<PrimitiveArray>();
268
269                DictArray::try_new(codes.into_array(), values.into_array())
270                    .vortex_unwrap()
271                    .into_array()
272            })
273            .collect::<ChunkedArray>()
274            .into_array()
275    }
276
277    #[test]
278    fn test_dict_array_from_primitive_chunks() {
279        let len = 2;
280        let chunk_count = 2;
281        let array = make_dict_primitive_chunks::<u64, u64>(len, 2, chunk_count);
282
283        let mut builder = builder_with_capacity(
284            &DType::Primitive(PType::U64, NonNullable),
285            len * chunk_count,
286        );
287        array
288            .clone()
289            .append_to_builder(builder.as_mut())
290            .vortex_unwrap();
291
292        let into_prim = array.to_primitive().unwrap();
293        let prim_into = builder.finish().to_primitive().unwrap();
294
295        assert_eq!(into_prim.as_slice::<u64>(), prim_into.as_slice::<u64>());
296        assert_eq!(
297            into_prim.validity_mask().unwrap().boolean_buffer(),
298            prim_into.validity_mask().unwrap().boolean_buffer()
299        )
300    }
301}
vortex_dict/array.rs

vortex_dict/
array.rs
