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!(
122                "Failed to get index {} from DictArray codes",
123                index
124            ))
125        })?;
126
127        if scalar.is_null() {
128            return Ok(false);
129        };
130        let values_index: usize = scalar
131            .as_ref()
132            .try_into()
133            .vortex_expect("Failed to convert dictionary code to usize");
134        array.values().is_valid(values_index)
135    }
136
137    fn all_valid(array: &DictArray) -> VortexResult<bool> {
138        Ok(array.codes().all_valid()? && array.values().all_valid()?)
139    }
140
141    fn all_invalid(array: &DictArray) -> VortexResult<bool> {
142        Ok(array.codes().all_invalid()? || array.values().all_invalid()?)
143    }
144
145    fn validity_mask(array: &DictArray) -> VortexResult<Mask> {
146        let codes_validity = array.codes().validity_mask()?;
147        match codes_validity.boolean_buffer() {
148            AllOr::All => {
149                let primitive_codes = array.codes().to_primitive()?;
150                let values_mask = array.values().validity_mask()?;
151                let is_valid_buffer = match_each_integer_ptype!(primitive_codes.ptype(), |$P| {
152                    let codes_slice = primitive_codes.as_slice::<$P>();
153                    BooleanBuffer::collect_bool(array.len(), |idx| {
154                       values_mask.value(codes_slice[idx] as usize)
155                    })
156                });
157                Ok(Mask::from_buffer(is_valid_buffer))
158            }
159            AllOr::None => Ok(Mask::AllFalse(array.len())),
160            AllOr::Some(validity_buff) => {
161                let primitive_codes = array.codes().to_primitive()?;
162                let values_mask = array.values().validity_mask()?;
163                let is_valid_buffer = match_each_integer_ptype!(primitive_codes.ptype(), |$P| {
164                    let codes_slice = primitive_codes.as_slice::<$P>();
165                    BooleanBuffer::collect_bool(array.len(), |idx| {
166                       validity_buff.value(idx) && values_mask.value(codes_slice[idx] as usize)
167                    })
168                });
169                Ok(Mask::from_buffer(is_valid_buffer))
170            }
171        }
172    }
173}
174
175#[cfg(test)]
176mod test {
177    use arrow_buffer::BooleanBuffer;
178    use rand::distr::{Distribution, StandardUniform};
179    use rand::prelude::StdRng;
180    use rand::{Rng, SeedableRng};
181    use vortex_array::arrays::{ChunkedArray, PrimitiveArray};
182    use vortex_array::builders::builder_with_capacity;
183    use vortex_array::validity::Validity;
184    use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
185    use vortex_buffer::buffer;
186    use vortex_dtype::Nullability::NonNullable;
187    use vortex_dtype::{DType, NativePType, PType};
188    use vortex_error::{VortexExpect, VortexUnwrap, vortex_panic};
189    use vortex_mask::AllOr;
190
191    use crate::DictArray;
192
193    #[test]
194    fn nullable_codes_validity() {
195        let dict = DictArray::try_new(
196            PrimitiveArray::new(
197                buffer![0u32, 1, 2, 2, 1],
198                Validity::from(BooleanBuffer::from(vec![true, false, true, false, true])),
199            )
200            .into_array(),
201            PrimitiveArray::new(buffer![3, 6, 9], Validity::AllValid).into_array(),
202        )
203        .unwrap();
204        let mask = dict.validity_mask().unwrap();
205        let AllOr::Some(indices) = mask.indices() else {
206            vortex_panic!("Expected indices from mask")
207        };
208        assert_eq!(indices, [0, 2, 4]);
209    }
210
211    #[test]
212    fn nullable_values_validity() {
213        let dict = DictArray::try_new(
214            buffer![0u32, 1, 2, 2, 1].into_array(),
215            PrimitiveArray::new(
216                buffer![3, 6, 9],
217                Validity::from(BooleanBuffer::from(vec![true, false, false])),
218            )
219            .into_array(),
220        )
221        .unwrap();
222        let mask = dict.validity_mask().unwrap();
223        let AllOr::Some(indices) = mask.indices() else {
224            vortex_panic!("Expected indices from mask")
225        };
226        assert_eq!(indices, [0]);
227    }
228
229    #[test]
230    fn nullable_codes_and_values() {
231        let dict = DictArray::try_new(
232            PrimitiveArray::new(
233                buffer![0u32, 1, 2, 2, 1],
234                Validity::from(BooleanBuffer::from(vec![true, false, true, false, true])),
235            )
236            .into_array(),
237            PrimitiveArray::new(
238                buffer![3, 6, 9],
239                Validity::from(BooleanBuffer::from(vec![false, true, true])),
240            )
241            .into_array(),
242        )
243        .unwrap();
244        let mask = dict.validity_mask().unwrap();
245        let AllOr::Some(indices) = mask.indices() else {
246            vortex_panic!("Expected indices from mask")
247        };
248        assert_eq!(indices, [2, 4]);
249    }
250
251    fn make_dict_primitive_chunks<T: NativePType, U: NativePType>(
252        len: usize,
253        unique_values: usize,
254        chunk_count: usize,
255    ) -> ArrayRef
256    where
257        StandardUniform: Distribution<T>,
258    {
259        let mut rng = StdRng::seed_from_u64(0);
260
261        (0..chunk_count)
262            .map(|_| {
263                let values = (0..unique_values)
264                    .map(|_| rng.random::<T>())
265                    .collect::<PrimitiveArray>();
266                let codes = (0..len)
267                    .map(|_| {
268                        U::from(rng.random_range(0..unique_values)).vortex_expect("valid value")
269                    })
270                    .collect::<PrimitiveArray>();
271
272                DictArray::try_new(codes.into_array(), values.into_array())
273                    .vortex_unwrap()
274                    .into_array()
275            })
276            .collect::<ChunkedArray>()
277            .into_array()
278    }
279
280    #[test]
281    fn test_dict_array_from_primitive_chunks() {
282        let len = 2;
283        let chunk_count = 2;
284        let array = make_dict_primitive_chunks::<u64, u64>(len, 2, chunk_count);
285
286        let mut builder = builder_with_capacity(
287            &DType::Primitive(PType::U64, NonNullable),
288            len * chunk_count,
289        );
290        array
291            .clone()
292            .append_to_builder(builder.as_mut())
293            .vortex_unwrap();
294
295        let into_prim = array.to_primitive().unwrap();
296        let prim_into = builder.finish().to_primitive().unwrap();
297
298        assert_eq!(into_prim.as_slice::<u64>(), prim_into.as_slice::<u64>());
299        assert_eq!(
300            into_prim.validity_mask().unwrap().boolean_buffer(),
301            prim_into.validity_mask().unwrap().boolean_buffer()
302        )
303    }
304}
vortex_dict/array.rs

vortex_dict/
array.rs
