vortex_array/builders/
lazy_validity_builder.rs

1use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer};
2use vortex_dtype::Nullability;
3use vortex_dtype::Nullability::{NonNullable, Nullable};
4use vortex_error::{VortexExpect, vortex_panic};
5use vortex_mask::Mask;
6
7use crate::validity::Validity;
8
9/// This is borrowed from arrow's null buffer builder, however we expose a `append_buffer`
10/// method to append a boolean buffer directly.
11pub struct LazyNullBufferBuilder {
12    inner: Option<BooleanBufferBuilder>,
13    len: usize,
14    capacity: usize,
15}
16
17impl LazyNullBufferBuilder {
18    /// Creates a new empty builder.
19    /// `capacity` is the number of bits in the null buffer.
20    pub fn new(capacity: usize) -> Self {
21        Self {
22            inner: None,
23            len: 0,
24            capacity,
25        }
26    }
27
28    #[inline]
29    pub fn append_n_non_nulls(&mut self, n: usize) {
30        if let Some(buf) = self.inner.as_mut() {
31            buf.append_n(n, true)
32        } else {
33            self.len += n;
34        }
35    }
36
37    #[inline]
38    pub fn append_non_null(&mut self) {
39        if let Some(buf) = self.inner.as_mut() {
40            buf.append(true)
41        } else {
42            self.len += 1;
43        }
44    }
45
46    #[inline]
47    pub fn append_n_nulls(&mut self, n: usize) {
48        self.materialize_if_needed();
49        self.inner
50            .as_mut()
51            .vortex_expect("cannot append null to non-nullable builder")
52            .append_n(n, false);
53    }
54
55    #[inline]
56    pub fn append_null(&mut self) {
57        self.materialize_if_needed();
58        self.inner
59            .as_mut()
60            .vortex_expect("cannot append null to non-nullable builder")
61            .append(false);
62    }
63
64    #[inline]
65    pub fn append(&mut self, not_null: bool) {
66        if not_null {
67            self.append_non_null()
68        } else {
69            self.append_null()
70        }
71    }
72
73    #[inline]
74    pub fn append_buffer(&mut self, bool_buffer: &BooleanBuffer) {
75        self.materialize_if_needed();
76        self.inner
77            .as_mut()
78            .vortex_expect("buffer just materialized")
79            .append_buffer(bool_buffer);
80    }
81
82    pub fn append_validity_mask(&mut self, validity_mask: Mask) {
83        match validity_mask {
84            Mask::AllTrue(len) => self.append_n_non_nulls(len),
85            Mask::AllFalse(len) => self.append_n_nulls(len),
86            Mask::Values(is_valid) => self.append_buffer(is_valid.boolean_buffer()),
87        }
88    }
89
90    pub fn set_bit(&mut self, index: usize, v: bool) {
91        self.materialize_if_needed();
92        self.inner
93            .as_mut()
94            .vortex_expect("buffer just materialized")
95            .set_bit(index, v);
96    }
97
98    pub fn len(&self) -> usize {
99        // self.len is the length of the builder if the inner buffer is not materialized
100        self.inner.as_ref().map(|i| i.len()).unwrap_or(self.len)
101    }
102
103    pub fn truncate(&mut self, len: usize) {
104        if let Some(b) = self.inner.as_mut() {
105            b.truncate(len)
106        }
107        self.len = len;
108    }
109
110    pub fn reserve(&mut self, n: usize) {
111        self.materialize_if_needed();
112        self.inner
113            .as_mut()
114            .vortex_expect("buffer just materialized")
115            .reserve(n);
116    }
117
118    pub fn finish(&mut self) -> Option<NullBuffer> {
119        self.len = 0;
120        Some(NullBuffer::new(self.inner.take()?.finish()))
121    }
122
123    pub fn finish_with_nullability(&mut self, nullability: Nullability) -> Validity {
124        let nulls = self.finish();
125        match (nullability, nulls) {
126            (NonNullable, None) => Validity::NonNullable,
127            (Nullable, None) => Validity::AllValid,
128            (Nullable, Some(arr)) => Validity::from(arr),
129            _ => vortex_panic!("Invalid nullability/nulls combination"),
130        }
131    }
132
133    pub fn ensure_capacity(&mut self, capacity: usize) {
134        if self.inner.is_none() {
135            self.capacity = capacity;
136        } else {
137            let inner = self
138                .inner
139                .as_mut()
140                .vortex_expect("buffer just materialized");
141            if capacity < inner.capacity() {
142                inner.reserve(capacity - inner.len());
143            }
144        }
145    }
146
147    #[inline]
148    fn materialize_if_needed(&mut self) {
149        if self.inner.is_none() {
150            self.materialize()
151        }
152    }
153
154    // This only happens once per builder
155    #[cold]
156    #[inline(never)]
157    fn materialize(&mut self) {
158        if self.inner.is_none() {
159            let mut b = BooleanBufferBuilder::new(self.len.max(self.capacity));
160            b.append_n(self.len, true);
161            self.inner = Some(b);
162        }
163    }
164}
vortex_array/builders/lazy_validity_builder.rs

vortex_array/builders/
lazy_validity_builder.rs
