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

vortex_array/arrays/dict/vtable/
mod.rs
