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

vortex_array/arrays/dict/vtable/
mod.rs
