Skip to main content

vortex_array/arrays/variant/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod operations;
5mod validity;
6
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_panic;
11use vortex_session::VortexSession;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::ExecutionResult;
16use crate::array::Array;
17use crate::array::ArrayId;
18use crate::array::ArrayView;
19use crate::array::EmptyArrayData;
20use crate::array::VTable;
21use crate::arrays::variant::SLOT_NAMES;
22use crate::buffer::BufferHandle;
23use crate::dtype::DType;
24use crate::serde::ArrayChildren;
25
26/// A [`Variant`]-encoded Vortex array.
27pub type VariantArray = Array<Variant>;
28
29#[derive(Clone, Debug)]
30pub struct Variant;
31
32impl Variant {
33    pub const ID: ArrayId = ArrayId::new_ref("vortex.variant");
34}
35
36impl VTable for Variant {
37    type ArrayData = EmptyArrayData;
38
39    type OperationsVTable = Self;
40
41    type ValidityVTable = Self;
42
43    fn id(&self) -> ArrayId {
44        Self::ID
45    }
46
47    fn validate(
48        &self,
49        _data: &Self::ArrayData,
50        dtype: &DType,
51        len: usize,
52        slots: &[Option<ArrayRef>],
53    ) -> VortexResult<()> {
54        vortex_ensure!(
55            slots[0].is_some(),
56            "VariantArray child slot must be present"
57        );
58        let child = slots[0]
59            .as_ref()
60            .vortex_expect("validated child slot presence");
61        vortex_ensure!(
62            matches!(dtype, DType::Variant(_)),
63            "Expected Variant DType, got {dtype}"
64        );
65        vortex_ensure!(
66            child.dtype() == dtype,
67            "VariantArray child dtype {} does not match outer dtype {}",
68            child.dtype(),
69            dtype
70        );
71        vortex_ensure!(
72            child.len() == len,
73            "VariantArray length {} does not match outer length {}",
74            child.len(),
75            len
76        );
77        Ok(())
78    }
79
80    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
81        0
82    }
83
84    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
85        vortex_panic!("VariantArray buffer index {idx} out of bounds")
86    }
87
88    fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
89        None
90    }
91
92    fn serialize(
93        _array: ArrayView<'_, Self>,
94        _session: &VortexSession,
95    ) -> VortexResult<Option<Vec<u8>>> {
96        Ok(Some(vec![]))
97    }
98
99    fn deserialize(
100        &self,
101        dtype: &DType,
102        len: usize,
103        metadata: &[u8],
104
105        _buffers: &[BufferHandle],
106        children: &dyn ArrayChildren,
107        _session: &VortexSession,
108    ) -> VortexResult<crate::array::ArrayParts<Self>> {
109        vortex_ensure!(
110            metadata.is_empty(),
111            "VariantArray expects empty metadata, got {} bytes",
112            metadata.len()
113        );
114        vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType");
115        vortex_ensure!(
116            children.len() == 1,
117            "Expected 1 child, got {}",
118            children.len()
119        );
120        // The child carries the nullability for the whole VariantArray.
121        let child = children.get(0, dtype, len)?;
122        Ok(
123            crate::array::ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData)
124                .with_slots(vec![Some(child)]),
125        )
126    }
127
128    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
129        match SLOT_NAMES.get(idx) {
130            Some(name) => (*name).to_string(),
131            None => vortex_panic!("VariantArray slot_name index {idx} out of bounds"),
132        }
133    }
134
135    fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
136        Ok(ExecutionResult::done(array))
137    }
138}
139
140#[cfg(test)]
141mod tests {}