vortex_array/arrays/varbin/vtable/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::DType;
5use vortex_dtype::Nullability;
6use vortex_dtype::PType;
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_err;
11
12use crate::ArrayRef;
13use crate::DeserializeMetadata;
14use crate::ExecutionCtx;
15use crate::IntoArray;
16use crate::ProstMetadata;
17use crate::SerializeMetadata;
18use crate::arrays::varbin::VarBinArray;
19use crate::buffer::BufferHandle;
20use crate::serde::ArrayChildren;
21use crate::validity::Validity;
22use crate::vtable;
23use crate::vtable::ArrayId;
24use crate::vtable::VTable;
25use crate::vtable::ValidityVTableFromValidityHelper;
26
27mod array;
28mod canonical;
29mod kernel;
30mod operations;
31mod validity;
32mod visitor;
33
34use canonical::varbin_to_canonical;
35use kernel::PARENT_KERNELS;
36use vortex_session::VortexSession;
37
38use crate::arrays::varbin::compute::rules::PARENT_RULES;
39
40vtable!(VarBin);
41
42#[derive(Clone, prost::Message)]
43pub struct VarBinMetadata {
44    #[prost(enumeration = "PType", tag = "1")]
45    pub(crate) offsets_ptype: i32,
46}
47
48impl VTable for VarBinVTable {
49    type Array = VarBinArray;
50
51    type Metadata = ProstMetadata<VarBinMetadata>;
52
53    type ArrayVTable = Self;
54    type OperationsVTable = Self;
55    type ValidityVTable = ValidityVTableFromValidityHelper;
56    type VisitorVTable = Self;
57
58    fn id(_array: &Self::Array) -> ArrayId {
59        Self::ID
60    }
61
62    fn metadata(array: &VarBinArray) -> VortexResult<Self::Metadata> {
63        Ok(ProstMetadata(VarBinMetadata {
64            offsets_ptype: PType::try_from(array.offsets().dtype())
65                .vortex_expect("Must be a valid PType") as i32,
66        }))
67    }
68
69    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
70        Ok(Some(metadata.serialize()))
71    }
72
73    fn deserialize(
74        bytes: &[u8],
75        _dtype: &DType,
76        _len: usize,
77        _buffers: &[BufferHandle],
78        _session: &VortexSession,
79    ) -> VortexResult<Self::Metadata> {
80        Ok(ProstMetadata(ProstMetadata::<VarBinMetadata>::deserialize(
81            bytes,
82        )?))
83    }
84
85    fn build(
86        dtype: &DType,
87        len: usize,
88        metadata: &Self::Metadata,
89        buffers: &[BufferHandle],
90        children: &dyn ArrayChildren,
91    ) -> VortexResult<VarBinArray> {
92        let validity = if children.len() == 1 {
93            Validity::from(dtype.nullability())
94        } else if children.len() == 2 {
95            let validity = children.get(1, &Validity::DTYPE, len)?;
96            Validity::Array(validity)
97        } else {
98            vortex_bail!("Expected 1 or 2 children, got {}", children.len());
99        };
100
101        let offsets = children.get(
102            0,
103            &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable),
104            len + 1,
105        )?;
106
107        if buffers.len() != 1 {
108            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
109        }
110        let bytes = buffers[0].clone().try_to_host_sync()?;
111
112        VarBinArray::try_new(offsets, bytes, dtype.clone(), validity)
113    }
114
115    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
116        match children.len() {
117            1 => {
118                let [offsets]: [ArrayRef; 1] = children
119                    .try_into()
120                    .map_err(|_| vortex_err!("Failed to convert children to array"))?;
121                array.offsets = offsets;
122            }
123            2 => {
124                let [offsets, validity]: [ArrayRef; 2] = children
125                    .try_into()
126                    .map_err(|_| vortex_err!("Failed to convert children to array"))?;
127                array.offsets = offsets;
128                array.validity = Validity::Array(validity);
129            }
130            _ => vortex_bail!(
131                "VarBinArray expects 1 or 2 children (offsets, validity?), got {}",
132                children.len()
133            ),
134        }
135        Ok(())
136    }
137
138    fn reduce_parent(
139        array: &Self::Array,
140        parent: &ArrayRef,
141        child_idx: usize,
142    ) -> VortexResult<Option<ArrayRef>> {
143        PARENT_RULES.evaluate(array, parent, child_idx)
144    }
145
146    fn execute_parent(
147        array: &Self::Array,
148        parent: &ArrayRef,
149        child_idx: usize,
150        ctx: &mut ExecutionCtx,
151    ) -> VortexResult<Option<ArrayRef>> {
152        PARENT_KERNELS.execute(array, parent, child_idx, ctx)
153    }
154
155    fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
156        Ok(varbin_to_canonical(array, ctx)?.into_array())
157    }
158}
159
160#[derive(Debug)]
161pub struct VarBinVTable;
162
163impl VarBinVTable {
164    pub const ID: ArrayId = ArrayId::new_ref("vortex.varbin");
165}
vortex_array/arrays/varbin/vtable/mod.rs

vortex_array/arrays/varbin/vtable/
mod.rs
