vortex_array/arrays/dict/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hasher;
5
6use kernel::PARENT_KERNELS;
7use prost::Message;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_ensure;
11use vortex_error::vortex_err;
12use vortex_error::vortex_panic;
13use vortex_session::VortexSession;
14use vortex_session::registry::CachedId;
15
16use super::DictData;
17use super::DictMetadata;
18use super::DictOwnedExt;
19use super::DictParts;
20use super::array::DictSlots;
21use super::array::DictSlotsView;
22use crate::AnyCanonical;
23use crate::ArrayEq;
24use crate::ArrayHash;
25use crate::ArrayRef;
26use crate::Canonical;
27use crate::Precision;
28use crate::array::Array;
29use crate::array::ArrayId;
30use crate::array::ArrayParts;
31use crate::array::ArrayView;
32use crate::array::VTable;
33use crate::arrays::ConstantArray;
34use crate::arrays::Primitive;
35use crate::arrays::dict::DictArrayExt;
36use crate::arrays::dict::DictArraySlotsExt;
37use crate::arrays::dict::compute::rules::PARENT_RULES;
38use crate::arrays::dict::execute::take_canonical;
39use crate::buffer::BufferHandle;
40use crate::dtype::DType;
41use crate::dtype::Nullability;
42use crate::dtype::PType;
43use crate::executor::ExecutionCtx;
44use crate::executor::ExecutionResult;
45use crate::require_child;
46use crate::scalar::Scalar;
47use crate::serde::ArrayChildren;
48use crate::validity::Validity;
49
50mod kernel;
51mod operations;
52mod validity;
53
54/// A [`Dict`]-encoded Vortex array.
55pub type DictArray = Array<Dict>;
56
57#[derive(Clone, Debug)]
58pub struct Dict;
59
60impl ArrayHash for DictData {
61    fn array_hash<H: Hasher>(&self, _state: &mut H, _precision: Precision) {}
62}
63
64impl ArrayEq for DictData {
65    fn array_eq(&self, _other: &Self, _precision: Precision) -> bool {
66        true
67    }
68}
69
70impl VTable for Dict {
71    type ArrayData = DictData;
72
73    type OperationsVTable = Self;
74    type ValidityVTable = Self;
75
76    fn id(&self) -> ArrayId {
77        static ID: CachedId = CachedId::new("vortex.dict");
78        *ID
79    }
80
81    fn validate(
82        &self,
83        _data: &DictData,
84        dtype: &DType,
85        len: usize,
86        slots: &[Option<ArrayRef>],
87    ) -> VortexResult<()> {
88        let view = DictSlotsView::from_slots(slots);
89        let codes = view.codes;
90        let values = view.values;
91        vortex_ensure!(codes.len() == len, "DictArray codes length mismatch");
92        vortex_ensure!(
93            values
94                .dtype()
95                .union_nullability(codes.dtype().nullability())
96                == *dtype,
97            "DictArray dtype does not match codes/values dtype"
98        );
99        Ok(())
100    }
101
102    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
103        0
104    }
105
106    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
107        vortex_panic!("DictArray buffer index {idx} out of bounds")
108    }
109
110    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
111        None
112    }
113
114    fn serialize(
115        array: ArrayView<'_, Self>,
116        _session: &VortexSession,
117    ) -> VortexResult<Option<Vec<u8>>> {
118        Ok(Some(
119            DictMetadata {
120                codes_ptype: PType::try_from(array.codes().dtype())? as i32,
121                values_len: u32::try_from(array.values().len()).map_err(|_| {
122                    vortex_err!(
123                        "Dictionary values size {} overflowed u32",
124                        array.values().len()
125                    )
126                })?,
127                is_nullable_codes: Some(array.codes().dtype().is_nullable()),
128                all_values_referenced: Some(array.has_all_values_referenced()),
129            }
130            .encode_to_vec(),
131        ))
132    }
133
134    fn deserialize(
135        &self,
136        dtype: &DType,
137        len: usize,
138        metadata: &[u8],
139        _buffers: &[BufferHandle],
140        children: &dyn ArrayChildren,
141        _session: &VortexSession,
142    ) -> VortexResult<ArrayParts<Self>> {
143        let metadata = DictMetadata::decode(metadata)?;
144        if children.len() != 2 {
145            vortex_bail!(
146                "Expected 2 children for dict encoding, found {}",
147                children.len()
148            )
149        }
150        let codes_nullable = metadata
151            .is_nullable_codes
152            .map(Nullability::from)
153            // If no `is_nullable_codes` metadata use the nullability of the values
154            // (and whole array) as before.
155            .unwrap_or_else(|| dtype.nullability());
156        let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable);
157        let codes = children.get(0, &codes_dtype, len)?;
158        let values = children.get(1, dtype, metadata.values_len as usize)?;
159        let all_values_referenced = metadata.all_values_referenced.unwrap_or(false);
160
161        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, unsafe {
162            DictData::new_unchecked().set_all_values_referenced(all_values_referenced)
163        })
164        .with_slots(vec![Some(codes), Some(values)]))
165    }
166
167    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
168        DictSlots::NAMES[idx].to_string()
169    }
170
171    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
172        if array.is_empty() {
173            let result_dtype = array
174                .dtype()
175                .union_nullability(array.codes().dtype().nullability());
176            return Ok(ExecutionResult::done(Canonical::empty(&result_dtype)));
177        }
178
179        let array = require_child!(array, array.codes(), DictSlots::CODES => Primitive);
180
181        if matches!(array.codes().validity()?, Validity::AllInvalid) {
182            return Ok(ExecutionResult::done(ConstantArray::new(
183                Scalar::null(array.dtype().as_nullable()),
184                array.codes().len(),
185            )));
186        }
187
188        let array = require_child!(array, array.values(), DictSlots::VALUES => AnyCanonical);
189
190        let DictParts { values, codes, .. } = array.into_parts();
191
192        Ok(ExecutionResult::done(take_canonical(
193            values.as_::<AnyCanonical>(),
194            &codes.downcast::<Primitive>(),
195            ctx,
196        )?))
197    }
198
199    fn reduce_parent(
200        array: ArrayView<'_, Self>,
201        parent: &ArrayRef,
202        child_idx: usize,
203    ) -> VortexResult<Option<ArrayRef>> {
204        PARENT_RULES.evaluate(array, parent, child_idx)
205    }
206
207    fn execute_parent(
208        array: ArrayView<'_, Self>,
209        parent: &ArrayRef,
210        child_idx: usize,
211        ctx: &mut ExecutionCtx,
212    ) -> VortexResult<Option<ArrayRef>> {
213        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
214    }
215}
vortex_array/arrays/dict/vtable/mod.rs

vortex_array/arrays/dict/vtable/
mod.rs
