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::ProstMetadata;
15use crate::SerializeMetadata;
16use crate::arrays::varbin::VarBinArray;
17use crate::buffer::BufferHandle;
18use crate::serde::ArrayChildren;
19use crate::validity::Validity;
20use crate::vtable;
21use crate::vtable::ArrayId;
22use crate::vtable::ArrayVTable;
23use crate::vtable::ArrayVTableExt;
24use crate::vtable::NotSupported;
25use crate::vtable::VTable;
26use crate::vtable::ValidityVTableFromValidityHelper;
27
28mod array;
29mod canonical;
30mod operations;
31mod validity;
32mod visitor;
33
34vtable!(VarBin);
35
36#[derive(Clone, prost::Message)]
37pub struct VarBinMetadata {
38    #[prost(enumeration = "PType", tag = "1")]
39    pub(crate) offsets_ptype: i32,
40}
41
42impl VTable for VarBinVTable {
43    type Array = VarBinArray;
44
45    type Metadata = ProstMetadata<VarBinMetadata>;
46
47    type ArrayVTable = Self;
48    type CanonicalVTable = Self;
49    type OperationsVTable = Self;
50    type ValidityVTable = ValidityVTableFromValidityHelper;
51    type VisitorVTable = Self;
52    type ComputeVTable = NotSupported;
53    type EncodeVTable = NotSupported;
54
55    fn id(&self) -> ArrayId {
56        ArrayId::new_ref("vortex.varbin")
57    }
58
59    fn encoding(_array: &Self::Array) -> ArrayVTable {
60        VarBinVTable.as_vtable()
61    }
62
63    fn metadata(array: &VarBinArray) -> VortexResult<Self::Metadata> {
64        Ok(ProstMetadata(VarBinMetadata {
65            offsets_ptype: PType::try_from(array.offsets().dtype())
66                .vortex_expect("Must be a valid PType") as i32,
67        }))
68    }
69
70    fn serialize(metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
71        Ok(Some(metadata.serialize()))
72    }
73
74    fn deserialize(bytes: &[u8]) -> VortexResult<Self::Metadata> {
75        Ok(ProstMetadata(ProstMetadata::<VarBinMetadata>::deserialize(
76            bytes,
77        )?))
78    }
79
80    fn build(
81        &self,
82        dtype: &DType,
83        len: usize,
84        metadata: &Self::Metadata,
85        buffers: &[BufferHandle],
86        children: &dyn ArrayChildren,
87    ) -> VortexResult<VarBinArray> {
88        let validity = if children.len() == 1 {
89            Validity::from(dtype.nullability())
90        } else if children.len() == 2 {
91            let validity = children.get(1, &Validity::DTYPE, len)?;
92            Validity::Array(validity)
93        } else {
94            vortex_bail!("Expected 1 or 2 children, got {}", children.len());
95        };
96
97        let offsets = children.get(
98            0,
99            &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable),
100            len + 1,
101        )?;
102
103        if buffers.len() != 1 {
104            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
105        }
106        let bytes = buffers[0].clone().try_to_bytes()?;
107
108        VarBinArray::try_new(offsets, bytes, dtype.clone(), validity)
109    }
110
111    fn with_children(array: &mut Self::Array, children: Vec<ArrayRef>) -> VortexResult<()> {
112        match children.len() {
113            1 => {
114                let [offsets]: [ArrayRef; 1] = children
115                    .try_into()
116                    .map_err(|_| vortex_err!("Failed to convert children to array"))?;
117                array.offsets = offsets;
118            }
119            2 => {
120                let [offsets, validity]: [ArrayRef; 2] = children
121                    .try_into()
122                    .map_err(|_| vortex_err!("Failed to convert children to array"))?;
123                array.offsets = offsets;
124                array.validity = Validity::Array(validity);
125            }
126            _ => vortex_bail!(
127                "VarBinArray expects 1 or 2 children (offsets, validity?), got {}",
128                children.len()
129            ),
130        }
131        Ok(())
132    }
133}
134
135#[derive(Debug)]
136pub struct VarBinVTable;