vortex_array/arrays/varbin/
builder.rs

1use arrow_buffer::NullBufferBuilder;
2use num_traits::{AsPrimitive, PrimInt};
3use vortex_buffer::BufferMut;
4use vortex_dtype::{DType, NativePType};
5use vortex_error::{VortexExpect as _, vortex_panic};
6
7use crate::IntoArray;
8use crate::arrays::primitive::PrimitiveArray;
9use crate::arrays::varbin::VarBinArray;
10use crate::validity::Validity;
11
12pub struct VarBinBuilder<O: NativePType> {
13    offsets: BufferMut<O>,
14    data: BufferMut<u8>,
15    validity: NullBufferBuilder,
16}
17
18impl<O: NativePType + PrimInt> Default for VarBinBuilder<O> {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl<O: NativePType + PrimInt> VarBinBuilder<O> {
25    pub fn new() -> Self {
26        Self::with_capacity(0)
27    }
28
29    pub fn with_capacity(len: usize) -> Self {
30        let mut offsets = BufferMut::with_capacity(len + 1);
31        offsets.push(O::zero());
32        Self {
33            offsets,
34            data: BufferMut::empty(),
35            validity: NullBufferBuilder::new(len),
36        }
37    }
38
39    #[inline]
40    pub fn append(&mut self, value: Option<&[u8]>) {
41        match value {
42            Some(v) => self.append_value(v),
43            None => self.append_null(),
44        }
45    }
46
47    #[inline]
48    pub fn append_value(&mut self, value: impl AsRef<[u8]>) {
49        let slice = value.as_ref();
50        self.offsets
51            .push(O::from(self.data.len() + slice.len()).unwrap_or_else(|| {
52                vortex_panic!(
53                    "Failed to convert sum of {} and {} to offset of type {}",
54                    self.data.len(),
55                    slice.len(),
56                    std::any::type_name::<O>()
57                )
58            }));
59        self.data.extend_from_slice(slice);
60        self.validity.append_non_null();
61    }
62
63    #[inline]
64    pub fn append_null(&mut self) {
65        self.offsets.push(self.offsets[self.offsets.len() - 1]);
66        self.validity.append_null();
67    }
68
69    #[inline]
70    pub fn append_n_nulls(&mut self, n: usize) {
71        self.offsets.push_n(self.offsets[self.offsets.len() - 1], n);
72        self.validity.append_n_nulls(n);
73    }
74
75    #[inline]
76    pub fn append_values(&mut self, values: &[u8], end_offsets: impl Iterator<Item = O>, num: usize)
77    where
78        O: 'static,
79        usize: AsPrimitive<O>,
80    {
81        self.offsets
82            .extend(end_offsets.map(|offset| offset + self.data.len().as_()));
83        self.data.extend_from_slice(values);
84        self.validity.append_n_non_nulls(num);
85    }
86
87    pub fn finish(mut self, dtype: DType) -> VarBinArray {
88        let offsets = PrimitiveArray::new(self.offsets.freeze(), Validity::NonNullable);
89        let nulls = self.validity.finish();
90
91        let validity = if dtype.is_nullable() {
92            nulls.map(Validity::from).unwrap_or(Validity::AllValid)
93        } else {
94            assert!(nulls.is_none(), "dtype and validity mismatch");
95            Validity::NonNullable
96        };
97
98        VarBinArray::try_new(offsets.into_array(), self.data.freeze(), dtype, validity)
99            .vortex_expect("Unexpected error while building VarBinArray")
100    }
101}
102
103#[cfg(test)]
104mod test {
105    use vortex_dtype::DType;
106    use vortex_dtype::Nullability::Nullable;
107    use vortex_scalar::Scalar;
108
109    use crate::arrays::varbin::builder::VarBinBuilder;
110
111    #[test]
112    fn test_builder() {
113        let mut builder = VarBinBuilder::<i32>::with_capacity(0);
114        builder.append(Some(b"hello"));
115        builder.append(None);
116        builder.append(Some(b"world"));
117        let array = builder.finish(DType::Utf8(Nullable));
118
119        assert_eq!(array.len(), 3);
120        assert_eq!(array.dtype().nullability(), Nullable);
121        assert_eq!(
122            array.scalar_at(0).unwrap(),
123            Scalar::utf8("hello".to_string(), Nullable)
124        );
125        assert!(array.scalar_at(1).unwrap().is_null());
126    }
127}
vortex_array/arrays/varbin/builder.rs

vortex_array/arrays/varbin/
builder.rs
