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 std::hash::Hasher;
8use std::sync::Arc;
9
10use vortex_error::VortexExpect;
11use vortex_error::VortexResult;
12use vortex_error::vortex_ensure;
13use vortex_error::vortex_panic;
14
15use crate::ArrayEq;
16use crate::ArrayHash;
17use crate::ArrayRef;
18use crate::EmptyMetadata;
19use crate::ExecutionCtx;
20use crate::ExecutionResult;
21use crate::Precision;
22use crate::arrays::VariantArray;
23use crate::buffer::BufferHandle;
24use crate::dtype::DType;
25use crate::serde::ArrayChildren;
26use crate::stats::StatsSetRef;
27use crate::vtable;
28use crate::vtable::ArrayId;
29use crate::vtable::VTable;
30
31vtable!(Variant);
32
33#[derive(Clone, Debug)]
34pub struct Variant;
35
36impl Variant {
37    pub const ID: ArrayId = ArrayId::new_ref("vortex.variant");
38}
39
40impl VTable for Variant {
41    type Array = VariantArray;
42
43    type Metadata = EmptyMetadata;
44
45    type OperationsVTable = Self;
46
47    type ValidityVTable = Self;
48
49    fn vtable(_array: &Self::Array) -> &Self {
50        &Variant
51    }
52
53    fn id(&self) -> ArrayId {
54        Self::ID
55    }
56
57    fn len(array: &Self::Array) -> usize {
58        array.child.len()
59    }
60
61    fn dtype(array: &Self::Array) -> &DType {
62        array.child.dtype()
63    }
64
65    fn stats(array: &Self::Array) -> StatsSetRef<'_> {
66        array.child.statistics()
67    }
68
69    fn array_hash<H: Hasher>(array: &Self::Array, state: &mut H, precision: Precision) {
70        array.child.array_hash(state, precision);
71    }
72
73    fn array_eq(array: &Self::Array, other: &Self::Array, precision: Precision) -> bool {
74        array.child.array_eq(&other.child, precision)
75    }
76
77    fn nbuffers(_array: &Self::Array) -> usize {
78        0
79    }
80
81    fn buffer(_array: &Self::Array, idx: usize) -> BufferHandle {
82        vortex_panic!("VariantArray buffer index {idx} out of bounds")
83    }
84
85    fn buffer_name(_array: &Self::Array, _idx: usize) -> Option<String> {
86        None
87    }
88
89    fn nchildren(_array: &Self::Array) -> usize {
90        1
91    }
92
93    fn child(array: &Self::Array, idx: usize) -> ArrayRef {
94        match idx {
95            0 => array.child.clone(),
96            _ => vortex_panic!("VariantArray child index {idx} out of bounds"),
97        }
98    }
99
100    fn child_name(_array: &Self::Array, idx: usize) -> String {
101        match idx {
102            0 => "child".to_string(),
103            _ => vortex_panic!("VariantArray child_name index {idx} out of bounds"),
104        }
105    }
106
107    fn metadata(_array: &Self::Array) -> VortexResult<Self::Metadata> {
108        Ok(EmptyMetadata)
109    }
110
111    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
112        Ok(Some(vec![]))
113    }
114
115    fn deserialize(
116        _bytes: &[u8],
117        _dtype: &DType,
118        _len: usize,
119        _buffers: &[BufferHandle],
120        _session: &vortex_session::VortexSession,
121    ) -> VortexResult<Self::Metadata> {
122        Ok(EmptyMetadata)
123    }
124
125    fn build(
126        dtype: &DType,
127        len: usize,
128        _metadata: &Self::Metadata,
129        _buffers: &[BufferHandle],
130        children: &dyn ArrayChildren,
131    ) -> VortexResult<Self::Array> {
132        vortex_ensure!(matches!(dtype, DType::Variant(_)), "Expected Variant DType");
133        vortex_ensure!(
134            children.len() == 1,
135            "Expected 1 child, got {}",
136            children.len()
137        );
138        // The child carries the nullability for the whole VariantArray.
139        let child = children.get(0, dtype, len)?;
140        Ok(VariantArray::new(child))
141    }
142
143    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
144        vortex_ensure!(
145            children.len() == 1,
146            "VariantArray expects exactly 1 child, got {}",
147            children.len()
148        );
149        array.child = children
150            .into_iter()
151            .next()
152            .vortex_expect("VariantArray must have 1 child");
153        Ok(())
154    }
155
156    fn execute(array: Arc<Self::Array>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
157        Ok(ExecutionResult::done_upcast::<Self>(array))
158    }
159}