vortex_array/arrays/varbin/
serde.rs

1use vortex_buffer::ByteBuffer;
2use vortex_dtype::{DType, Nullability, PType};
3use vortex_error::{VortexExpect, VortexResult, vortex_bail};
4
5use super::VarBinEncoding;
6use crate::arrays::{VarBinArray, VarBinVTable};
7use crate::serde::ArrayChildren;
8use crate::validity::Validity;
9use crate::vtable::{SerdeVTable, ValidityHelper, VisitorVTable};
10use crate::{Array, ArrayBufferVisitor, ArrayChildVisitor, ProstMetadata};
11
12#[derive(Clone, prost::Message)]
13pub struct VarBinMetadata {
14    #[prost(enumeration = "PType", tag = "1")]
15    pub(crate) offsets_ptype: i32,
16}
17
18impl SerdeVTable<VarBinVTable> for VarBinVTable {
19    type Metadata = ProstMetadata<VarBinMetadata>;
20
21    fn metadata(array: &VarBinArray) -> VortexResult<Option<Self::Metadata>> {
22        Ok(Some(ProstMetadata(VarBinMetadata {
23            offsets_ptype: PType::try_from(array.offsets().dtype())
24                .vortex_expect("Must be a valid PType") as i32,
25        })))
26    }
27
28    fn build(
29        _encoding: &VarBinEncoding,
30        dtype: &DType,
31        len: usize,
32        metadata: &VarBinMetadata,
33        buffers: &[ByteBuffer],
34        children: &dyn ArrayChildren,
35    ) -> VortexResult<VarBinArray> {
36        let validity = if children.len() == 1 {
37            Validity::from(dtype.nullability())
38        } else if children.len() == 2 {
39            let validity = children.get(1, &Validity::DTYPE, len)?;
40            Validity::Array(validity)
41        } else {
42            vortex_bail!("Expected 1 or 2 children, got {}", children.len());
43        };
44
45        let offsets = children.get(
46            0,
47            &DType::Primitive(metadata.offsets_ptype(), Nullability::NonNullable),
48            len + 1,
49        )?;
50
51        if buffers.len() != 1 {
52            vortex_bail!("Expected 1 buffer, got {}", buffers.len());
53        }
54        let bytes = buffers[0].clone();
55
56        VarBinArray::try_new(offsets, bytes, dtype.clone(), validity)
57    }
58}
59
60impl VisitorVTable<VarBinVTable> for VarBinVTable {
61    fn visit_buffers(array: &VarBinArray, visitor: &mut dyn ArrayBufferVisitor) {
62        visitor.visit_buffer(array.bytes()); // TODO(ngates): sliced bytes?
63    }
64
65    fn visit_children(array: &VarBinArray, visitor: &mut dyn ArrayChildVisitor) {
66        visitor.visit_child("offsets", array.offsets());
67        visitor.visit_validity(array.validity(), array.len());
68    }
69}