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

vortex_array/arrays/dict/vtable/
mod.rs
