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

vortex_array/arrays/dict/vtable/
mod.rs
