arrow_buffer/builder/
null.rs

Help
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::{BooleanBufferBuilder, MutableBuffer, NullBuffer};
19
20/// Builder for creating the null bit buffer.
21///
22/// This builder only materializes the buffer when we append `false`.
23/// If you only append `true`s to the builder, what you get will be
24/// `None` when calling [`finish`](#method.finish).
25/// This optimization is **very** important for the performance.
26#[derive(Debug)]
27pub struct NullBufferBuilder {
28    bitmap_builder: Option<BooleanBufferBuilder>,
29    /// Store the length of the buffer before materializing.
30    len: usize,
31    capacity: usize,
32}
33
34impl NullBufferBuilder {
35    /// Creates a new empty builder.
36    /// `capacity` is the number of bits in the null buffer.
37    pub fn new(capacity: usize) -> Self {
38        Self {
39            bitmap_builder: None,
40            len: 0,
41            capacity,
42        }
43    }
44
45    /// Creates a new builder with given length.
46    pub fn new_with_len(len: usize) -> Self {
47        Self {
48            bitmap_builder: None,
49            len,
50            capacity: len,
51        }
52    }
53
54    /// Creates a new builder from a `MutableBuffer`.
55    pub fn new_from_buffer(buffer: MutableBuffer, len: usize) -> Self {
56        let capacity = buffer.len() * 8;
57
58        assert!(len <= capacity);
59
60        let bitmap_builder = Some(BooleanBufferBuilder::new_from_buffer(buffer, len));
61        Self {
62            bitmap_builder,
63            len,
64            capacity,
65        }
66    }
67
68    /// Appends `n` `true`s into the builder
69    /// to indicate that these `n` items are not nulls.
70    #[inline]
71    pub fn append_n_non_nulls(&mut self, n: usize) {
72        if let Some(buf) = self.bitmap_builder.as_mut() {
73            buf.append_n(n, true)
74        } else {
75            self.len += n;
76        }
77    }
78
79    /// Appends a `true` into the builder
80    /// to indicate that this item is not null.
81    #[inline]
82    pub fn append_non_null(&mut self) {
83        if let Some(buf) = self.bitmap_builder.as_mut() {
84            buf.append(true)
85        } else {
86            self.len += 1;
87        }
88    }
89
90    /// Appends `n` `false`s into the builder
91    /// to indicate that these `n` items are nulls.
92    #[inline]
93    pub fn append_n_nulls(&mut self, n: usize) {
94        self.materialize_if_needed();
95        self.bitmap_builder.as_mut().unwrap().append_n(n, false);
96    }
97
98    /// Appends a `false` into the builder
99    /// to indicate that this item is null.
100    #[inline]
101    pub fn append_null(&mut self) {
102        self.materialize_if_needed();
103        self.bitmap_builder.as_mut().unwrap().append(false);
104    }
105
106    /// Appends a boolean value into the builder.
107    #[inline]
108    pub fn append(&mut self, not_null: bool) {
109        if not_null {
110            self.append_non_null()
111        } else {
112            self.append_null()
113        }
114    }
115
116    /// Appends a boolean slice into the builder
117    /// to indicate the validations of these items.
118    pub fn append_slice(&mut self, slice: &[bool]) {
119        if slice.iter().any(|v| !v) {
120            self.materialize_if_needed()
121        }
122        if let Some(buf) = self.bitmap_builder.as_mut() {
123            buf.append_slice(slice)
124        } else {
125            self.len += slice.len();
126        }
127    }
128
129    /// Builds the null buffer and resets the builder.
130    /// Returns `None` if the builder only contains `true`s.
131    pub fn finish(&mut self) -> Option<NullBuffer> {
132        self.len = 0;
133        Some(NullBuffer::new(self.bitmap_builder.take()?.finish()))
134    }
135
136    /// Builds the [NullBuffer] without resetting the builder.
137    pub fn finish_cloned(&self) -> Option<NullBuffer> {
138        let buffer = self.bitmap_builder.as_ref()?.finish_cloned();
139        Some(NullBuffer::new(buffer))
140    }
141
142    /// Returns the inner bitmap builder as slice
143    pub fn as_slice(&self) -> Option<&[u8]> {
144        Some(self.bitmap_builder.as_ref()?.as_slice())
145    }
146
147    fn materialize_if_needed(&mut self) {
148        if self.bitmap_builder.is_none() {
149            self.materialize()
150        }
151    }
152
153    #[cold]
154    fn materialize(&mut self) {
155        if self.bitmap_builder.is_none() {
156            let mut b = BooleanBufferBuilder::new(self.len.max(self.capacity));
157            b.append_n(self.len, true);
158            self.bitmap_builder = Some(b);
159        }
160    }
161
162    /// Return a mutable reference to the inner bitmap slice.
163    pub fn as_slice_mut(&mut self) -> Option<&mut [u8]> {
164        self.bitmap_builder.as_mut().map(|b| b.as_slice_mut())
165    }
166
167    /// Return the allocated size of this builder, in bytes, useful for memory accounting.
168    pub fn allocated_size(&self) -> usize {
169        self.bitmap_builder
170            .as_ref()
171            .map(|b| b.capacity())
172            .unwrap_or(0)
173    }
174}
175
176impl NullBufferBuilder {
177    /// Return the number of bits in the buffer.
178    pub fn len(&self) -> usize {
179        self.bitmap_builder.as_ref().map_or(self.len, |b| b.len())
180    }
181
182    /// Check if the builder is empty.
183    pub fn is_empty(&self) -> bool {
184        self.len() == 0
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[test]
193    fn test_null_buffer_builder() {
194        let mut builder = NullBufferBuilder::new(0);
195        builder.append_null();
196        builder.append_non_null();
197        builder.append_n_nulls(2);
198        builder.append_n_non_nulls(2);
199        assert_eq!(6, builder.len());
200
201        let buf = builder.finish().unwrap();
202        assert_eq!(&[0b110010_u8], buf.validity());
203    }
204
205    #[test]
206    fn test_null_buffer_builder_all_nulls() {
207        let mut builder = NullBufferBuilder::new(0);
208        builder.append_null();
209        builder.append_n_nulls(2);
210        builder.append_slice(&[false, false, false]);
211        assert_eq!(6, builder.len());
212
213        let buf = builder.finish().unwrap();
214        assert_eq!(&[0b0_u8], buf.validity());
215    }
216
217    #[test]
218    fn test_null_buffer_builder_no_null() {
219        let mut builder = NullBufferBuilder::new(0);
220        builder.append_non_null();
221        builder.append_n_non_nulls(2);
222        builder.append_slice(&[true, true, true]);
223        assert_eq!(6, builder.len());
224
225        let buf = builder.finish();
226        assert!(buf.is_none());
227    }
228
229    #[test]
230    fn test_null_buffer_builder_reset() {
231        let mut builder = NullBufferBuilder::new(0);
232        builder.append_slice(&[true, false, true]);
233        builder.finish();
234        assert!(builder.is_empty());
235
236        builder.append_slice(&[true, true, true]);
237        assert!(builder.finish().is_none());
238        assert!(builder.is_empty());
239
240        builder.append_slice(&[true, true, false, true]);
241
242        let buf = builder.finish().unwrap();
243        assert_eq!(&[0b1011_u8], buf.validity());
244    }
245}
arrow_buffer/builder/null.rs

arrow_buffer/builder/
null.rs
