vortex_array/arrays/varbin/
serde.rs

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