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::Hash;
5
6use kernel::PARENT_KERNELS;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_err;
11use vortex_error::vortex_panic;
12use vortex_session::VortexSession;
13
14use super::DictArray;
15use super::DictMetadata;
16use super::take_canonical;
17use crate::Array;
18use crate::ArrayRef;
19use crate::Canonical;
20use crate::DeserializeMetadata;
21use crate::IntoArray;
22use crate::Precision;
23use crate::ProstMetadata;
24use crate::SerializeMetadata;
25use crate::arrays::ConstantArray;
26use crate::arrays::dict::compute::rules::PARENT_RULES;
27use crate::buffer::BufferHandle;
28use crate::dtype::DType;
29use crate::dtype::Nullability;
30use crate::dtype::PType;
31use crate::executor::ExecutionCtx;
32use crate::hash::ArrayEq;
33use crate::hash::ArrayHash;
34use crate::scalar::Scalar;
35use crate::serde::ArrayChildren;
36use crate::stats::StatsSetRef;
37use crate::vtable;
38use crate::vtable::ArrayId;
39use crate::vtable::VTable;
40mod kernel;
41mod operations;
42mod validity;
43
44vtable!(Dict);
45
46#[derive(Debug)]
47pub struct DictVTable;
48
49impl DictVTable {
50    pub const ID: ArrayId = ArrayId::new_ref("vortex.dict");
51}
52
53impl VTable for DictVTable {
54    type Array = DictArray;
55
56    type Metadata = ProstMetadata<DictMetadata>;
57    type OperationsVTable = Self;
58    type ValidityVTable = Self;
59
60    fn id(_array: &Self::Array) -> ArrayId {
61        Self::ID
62    }
63
64    fn len(array: &DictArray) -> usize {
65        array.codes.len()
66    }
67
68    fn dtype(array: &DictArray) -> &DType {
69        &array.dtype
70    }
71
72    fn stats(array: &DictArray) -> StatsSetRef<'_> {
73        array.stats_set.to_ref(array.as_ref())
74    }
75
76    fn array_hash<H: std::hash::Hasher>(array: &DictArray, state: &mut H, precision: Precision) {
77        array.dtype.hash(state);
78        array.codes.array_hash(state, precision);
79        array.values.array_hash(state, precision);
80    }
81
82    fn array_eq(array: &DictArray, other: &DictArray, precision: Precision) -> bool {
83        array.dtype == other.dtype
84            && array.codes.array_eq(&other.codes, precision)
85            && array.values.array_eq(&other.values, precision)
86    }
87
88    fn nbuffers(_array: &DictArray) -> usize {
89        0
90    }
91
92    fn buffer(_array: &DictArray, idx: usize) -> BufferHandle {
93        vortex_panic!("DictArray buffer index {idx} out of bounds")
94    }
95
96    fn buffer_name(_array: &DictArray, _idx: usize) -> Option<String> {
97        None
98    }
99
100    fn nchildren(_array: &DictArray) -> usize {
101        2
102    }
103
104    fn child(array: &DictArray, idx: usize) -> ArrayRef {
105        match idx {
106            0 => array.codes().clone(),
107            1 => array.values().clone(),
108            _ => vortex_panic!("DictArray child index {idx} out of bounds"),
109        }
110    }
111
112    fn child_name(_array: &DictArray, idx: usize) -> String {
113        match idx {
114            0 => "codes".to_string(),
115            1 => "values".to_string(),
116            _ => vortex_panic!("DictArray child_name index {idx} out of bounds"),
117        }
118    }
119
120    fn metadata(array: &DictArray) -> VortexResult<Self::Metadata> {
121        Ok(ProstMetadata(DictMetadata {
122            codes_ptype: PType::try_from(array.codes().dtype())? as i32,
123            values_len: u32::try_from(array.values().len()).map_err(|_| {
124                vortex_err!(
125                    "Dictionary values size {} overflowed u32",
126                    array.values().len()
127                )
128            })?,
129            is_nullable_codes: Some(array.codes().dtype().is_nullable()),
130            all_values_referenced: Some(array.all_values_referenced),
131        }))
132    }
133
134    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
135        Ok(Some(metadata.serialize()))
136    }
137
138    fn deserialize(
139        bytes: &[u8],
140        _dtype: &DType,
141        _len: usize,
142        _buffers: &[BufferHandle],
143        _session: &VortexSession,
144    ) -> VortexResult<Self::Metadata> {
145        let metadata = <Self::Metadata as DeserializeMetadata>::deserialize(bytes)?;
146        Ok(ProstMetadata(metadata))
147    }
148
149    fn build(
150        dtype: &DType,
151        len: usize,
152        metadata: &Self::Metadata,
153        _buffers: &[BufferHandle],
154        children: &dyn ArrayChildren,
155    ) -> VortexResult<DictArray> {
156        if children.len() != 2 {
157            vortex_bail!(
158                "Expected 2 children for dict encoding, found {}",
159                children.len()
160            )
161        }
162        let codes_nullable = metadata
163            .is_nullable_codes
164            .map(Nullability::from)
165            // If no `is_nullable_codes` metadata use the nullability of the values
166            // (and whole array) as before.
167            .unwrap_or_else(|| dtype.nullability());
168        let codes_dtype = DType::Primitive(metadata.codes_ptype(), codes_nullable);
169        let codes = children.get(0, &codes_dtype, len)?;
170        let values = children.get(1, dtype, metadata.values_len as usize)?;
171        let all_values_referenced = metadata.all_values_referenced.unwrap_or(false);
172
173        // SAFETY: We've validated the metadata and children.
174        Ok(unsafe {
175            DictArray::new_unchecked(codes, values).set_all_values_referenced(all_values_referenced)
176        })
177    }
178
179    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
180        vortex_ensure!(
181            children.len() == 2,
182            "DictArray expects exactly 2 children (codes, values), got {}",
183            children.len()
184        );
185        let [codes, values]: [ArrayRef; 2] = children
186            .try_into()
187            .map_err(|_| vortex_err!("Failed to convert children to array"))?;
188        array.codes = codes;
189        array.values = values;
190        Ok(())
191    }
192
193    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
194        if let Some(canonical) = execute_fast_path(array, ctx)? {
195            return Ok(canonical);
196        }
197
198        // TODO(joe): if the values are constant return a constant
199        let values = array.values().clone().execute::<Canonical>(ctx)?;
200        let codes = array
201            .codes()
202            .clone()
203            .execute::<Canonical>(ctx)?
204            .into_primitive();
205
206        // TODO(ngates): if indices are sorted and unique (strict-sorted), then we should delegate to
207        //  the filter function since they're typically optimised for this case.
208        // TODO(ngates): if indices min is quite high, we could slice self and offset the indices
209        //  such that canonicalize does less work.
210
211        Ok(take_canonical(values, &codes, ctx)?.into_array())
212    }
213
214    fn reduce_parent(
215        array: &Self::Array,
216        parent: &ArrayRef,
217        child_idx: usize,
218    ) -> VortexResult<Option<ArrayRef>> {
219        PARENT_RULES.evaluate(array, parent, child_idx)
220    }
221
222    fn execute_parent(
223        array: &Self::Array,
224        parent: &ArrayRef,
225        child_idx: usize,
226        ctx: &mut ExecutionCtx,
227    ) -> VortexResult<Option<ArrayRef>> {
228        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
229    }
230}
231
232/// Check for fast-path execution conditions.
233pub(super) fn execute_fast_path(
234    array: &DictArray,
235    _ctx: &mut ExecutionCtx,
236) -> VortexResult<Option<ArrayRef>> {
237    // Empty array - nothing to do
238    if array.is_empty() {
239        let result_dtype = array
240            .dtype()
241            .union_nullability(array.codes().dtype().nullability());
242        return Ok(Some(Canonical::empty(&result_dtype).into_array()));
243    }
244
245    // All codes are null - result is all nulls
246    if array.codes.all_invalid()? {
247        return Ok(Some(
248            ConstantArray::new(Scalar::null(array.dtype().as_nullable()), array.codes.len())
249                .into_array(),
250        ));
251    }
252
253    Ok(None)
254}
vortex_array/arrays/dict/vtable/mod.rs

vortex_array/arrays/dict/vtable/
mod.rs
