vortex_array/arrays/varbin/
builder.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::AsPrimitive;
5use vortex_buffer::BitBufferMut;
6use vortex_buffer::BufferMut;
7use vortex_error::vortex_panic;
8
9use crate::IntoArray;
10#[cfg(debug_assertions)]
11use crate::VortexSessionExecute;
12use crate::arrays::PrimitiveArray;
13use crate::arrays::VarBinArray;
14use crate::dtype::DType;
15use crate::dtype::IntegerPType;
16use crate::expr::stats::Precision;
17use crate::expr::stats::Stat;
18#[cfg(debug_assertions)]
19use crate::legacy_session;
20use crate::validity::Validity;
21
22pub struct VarBinBuilder<O: IntegerPType> {
23    offsets: BufferMut<O>,
24    data: BufferMut<u8>,
25    validity: BitBufferMut,
26}
27
28impl<O: IntegerPType> Default for VarBinBuilder<O> {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl<O: IntegerPType> VarBinBuilder<O> {
35    pub fn new() -> Self {
36        Self::with_capacity(0)
37    }
38
39    pub fn with_capacity(len: usize) -> Self {
40        let mut offsets = BufferMut::with_capacity(len + 1);
41        offsets.push(O::zero());
42        Self {
43            offsets,
44            data: BufferMut::empty(),
45            validity: BitBufferMut::with_capacity(len),
46        }
47    }
48
49    #[inline]
50    pub fn append(&mut self, value: Option<&[u8]>) {
51        match value {
52            Some(v) => self.append_value(v),
53            None => self.append_null(),
54        }
55    }
56
57    #[inline]
58    pub fn append_value(&mut self, value: impl AsRef<[u8]>) {
59        let slice = value.as_ref();
60        self.offsets
61            .push(O::from(self.data.len() + slice.len()).unwrap_or_else(|| {
62                vortex_panic!(
63                    "Failed to convert sum of {} and {} to offset of type {}",
64                    self.data.len(),
65                    slice.len(),
66                    std::any::type_name::<O>()
67                )
68            }));
69        self.data.extend_from_slice(slice);
70        self.validity.append_true();
71    }
72
73    #[inline]
74    pub fn append_null(&mut self) {
75        self.offsets.push(self.offsets[self.offsets.len() - 1]);
76        self.validity.append_false();
77    }
78
79    #[inline]
80    pub fn append_n_nulls(&mut self, n: usize) {
81        self.offsets.push_n(self.offsets[self.offsets.len() - 1], n);
82        self.validity.append_n(false, n);
83    }
84
85    #[inline]
86    pub fn append_values(&mut self, values: &[u8], end_offsets: impl Iterator<Item = O>, num: usize)
87    where
88        O: 'static,
89        usize: AsPrimitive<O>,
90    {
91        self.offsets
92            .extend(end_offsets.map(|offset| offset + self.data.len().as_()));
93        self.data.extend_from_slice(values);
94        self.validity.append_n(true, num);
95    }
96
97    #[allow(clippy::disallowed_methods)]
98    pub fn finish(self, dtype: DType) -> VarBinArray {
99        let offsets = PrimitiveArray::new(self.offsets.freeze(), Validity::NonNullable);
100        let nulls = self.validity.freeze();
101
102        let validity = Validity::from_bit_buffer(nulls, dtype.nullability());
103
104        // The builder guarantees offsets are monotonically increasing, so we can set
105        // this stat eagerly. This avoids an O(n) recomputation when the array is
106        // deserialized and VarBinArray::validate checks sortedness.
107        #[cfg(debug_assertions)]
108        {
109            let offsets_are_sorted = offsets
110                .statistics()
111                .compute_is_sorted(&mut legacy_session().create_execution_ctx())
112                .unwrap_or(false);
113            debug_assert!(offsets_are_sorted, "VarBinBuilder offsets must be sorted");
114        }
115        offsets
116            .statistics()
117            .set(Stat::IsSorted, Precision::Exact(true.into()));
118
119        // SAFETY: The builder maintains all invariants:
120        // - Offsets are monotonically increasing starting from 0 (guaranteed by builder logic).
121        // - Bytes buffer contains exactly the data referenced by offsets.
122        // - Validity matches the dtype nullability.
123        // - UTF-8 validity is ensured by the caller when using DType::Utf8.
124        unsafe {
125            VarBinArray::new_unchecked(offsets.into_array(), self.data.freeze(), dtype, validity)
126        }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use vortex_error::VortexResult;
133
134    use crate::VortexSessionExecute;
135    use crate::array_session;
136    use crate::arrays::varbin::VarBinArrayExt;
137    use crate::arrays::varbin::builder::VarBinBuilder;
138    use crate::dtype::DType;
139    use crate::dtype::Nullability::Nullable;
140    use crate::expr::stats::Precision;
141    use crate::expr::stats::Stat;
142    use crate::expr::stats::StatsProviderExt;
143    use crate::scalar::Scalar;
144
145    #[test]
146    fn test_builder() {
147        let mut builder = VarBinBuilder::<i32>::with_capacity(0);
148        builder.append(Some(b"hello"));
149        builder.append(None);
150        builder.append(Some(b"world"));
151        let array = builder.finish(DType::Utf8(Nullable));
152
153        assert_eq!(array.len(), 3);
154        assert_eq!(array.dtype().nullability(), Nullable);
155        assert_eq!(
156            array
157                .execute_scalar(0, &mut array_session().create_execution_ctx())
158                .unwrap(),
159            Scalar::utf8("hello".to_string(), Nullable)
160        );
161        assert!(
162            array
163                .execute_scalar(1, &mut array_session().create_execution_ctx())
164                .unwrap()
165                .is_null()
166        );
167    }
168
169    #[test]
170    fn offsets_have_is_sorted_stat() -> VortexResult<()> {
171        let mut builder = VarBinBuilder::<i32>::with_capacity(0);
172        builder.append_value(b"aaa");
173        builder.append_null();
174        builder.append_value(b"bbb");
175        let array = builder.finish(DType::Utf8(Nullable));
176
177        let is_sorted = array
178            .offsets()
179            .statistics()
180            .with_typed_stats_set(|s| s.get_as::<bool>(Stat::IsSorted));
181        assert_eq!(is_sorted, Precision::Exact(true));
182        Ok(())
183    }
184
185    #[test]
186    fn empty_builder_offsets_have_is_sorted_stat() -> VortexResult<()> {
187        let builder = VarBinBuilder::<i32>::new();
188        let array = builder.finish(DType::Utf8(Nullable));
189
190        let is_sorted = array
191            .offsets()
192            .statistics()
193            .with_typed_stats_set(|s| s.get_as::<bool>(Stat::IsSorted));
194        assert_eq!(is_sorted, Precision::Exact(true));
195        Ok(())
196    }
197}
vortex_array/arrays/varbin/builder.rs

vortex_array/arrays/varbin/
builder.rs
