arrow_buffer/buffer/
null.rs

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::bit_iterator::{BitIndexIterator, BitIterator, BitSliceIterator};
19use crate::buffer::BooleanBuffer;
20use crate::{Buffer, MutableBuffer};
21
22/// A [`BooleanBuffer`] used to encode validity (null values) for Arrow arrays
23///
24/// In the [Arrow specification], array validity is encoded in a packed bitmask with a
25/// `true` value indicating the corresponding slot is not null, and `false` indicating
26/// that it is null.
27///
28/// # See also
29/// * [`NullBufferBuilder`] for creating `NullBuffer`s
30///
31/// [Arrow specification]: https://arrow.apache.org/docs/format/Columnar.html#validity-bitmaps
32/// [`NullBufferBuilder`]: crate::NullBufferBuilder
33#[derive(Debug, Clone, Eq, PartialEq)]
34pub struct NullBuffer {
35    buffer: BooleanBuffer,
36    null_count: usize,
37}
38
39impl NullBuffer {
40    /// Create a new [`NullBuffer`] computing the null count
41    pub fn new(buffer: BooleanBuffer) -> Self {
42        let null_count = buffer.len() - buffer.count_set_bits();
43        Self { buffer, null_count }
44    }
45
46    /// Create a new [`NullBuffer`] of length `len` where all values are null
47    pub fn new_null(len: usize) -> Self {
48        Self {
49            buffer: BooleanBuffer::new_unset(len),
50            null_count: len,
51        }
52    }
53
54    /// Create a new [`NullBuffer`] of length `len` where all values are valid
55    ///
56    /// Note: it is more efficient to not set the null buffer if it is known to
57    /// be all valid (aka all values are not null)
58    pub fn new_valid(len: usize) -> Self {
59        Self {
60            buffer: BooleanBuffer::new_set(len),
61            null_count: 0,
62        }
63    }
64
65    /// Create a new [`NullBuffer`] with the provided `buffer` and `null_count`
66    ///
67    /// # Safety
68    ///
69    /// `buffer` must contain `null_count` `0` bits
70    pub unsafe fn new_unchecked(buffer: BooleanBuffer, null_count: usize) -> Self {
71        Self { buffer, null_count }
72    }
73
74    /// Computes the union of the nulls in two optional [`NullBuffer`]
75    ///
76    /// This is commonly used by binary operations where the result is NULL if either
77    /// of the input values is NULL. Handling the null mask separately in this way
78    /// can yield significant performance improvements over an iterator approach
79    pub fn union(lhs: Option<&NullBuffer>, rhs: Option<&NullBuffer>) -> Option<NullBuffer> {
80        match (lhs, rhs) {
81            (Some(lhs), Some(rhs)) if lhs.null_count() > 0 || rhs.null_count() > 0 => {
82                Some(Self::new(lhs.inner() & rhs.inner()))
83            }
84            (Some(n), None) | (None, Some(n)) if n.null_count() > 0 => Some(n.clone()),
85            (_, _) => None,
86        }
87    }
88
89    /// Computes the union of the nulls in multiple optional [`NullBuffer`]s
90    ///
91    /// See [`union`](Self::union)
92    pub fn union_many<'a>(
93        nulls: impl IntoIterator<Item = Option<&'a NullBuffer>>,
94    ) -> Option<NullBuffer> {
95        // Unwrap to BooleanBuffer because BitAndAssign is not implemented for NullBuffer
96        let mut buffers = nulls.into_iter().filter_map(|nb| match nb {
97            Some(nb) if nb.null_count > 0 => Some(nb.inner()),
98            _ => None,
99        });
100        let first = buffers.next()?;
101        let mut result = first.clone();
102        for buf in buffers {
103            result &= buf;
104        }
105        Some(Self::new(result))
106    }
107
108    /// Returns true if all nulls in `other` also exist in self
109    pub fn contains(&self, other: &NullBuffer) -> bool {
110        if other.null_count == 0 {
111            return true;
112        }
113        let lhs = self.inner().bit_chunks().iter_padded();
114        let rhs = other.inner().bit_chunks().iter_padded();
115        lhs.zip(rhs).all(|(l, r)| (l & !r) == 0)
116    }
117
118    /// Returns a new [`NullBuffer`] where each bit in the current null buffer
119    /// is repeated `count` times. This is useful for masking the nulls of
120    /// the child of a FixedSizeListArray based on its parent
121    ///
122    /// # Panics
123    ///
124    /// Panics if `self.len() * count` overflows `usize`
125    pub fn expand(&self, count: usize) -> Self {
126        let capacity = self.buffer.len().checked_mul(count).unwrap();
127        let mut buffer = MutableBuffer::new_null(capacity);
128
129        // Expand each bit within `null_mask` into `element_len`
130        // bits, constructing the implicit mask of the child elements
131        for i in 0..self.buffer.len() {
132            if self.is_null(i) {
133                continue;
134            }
135            for j in 0..count {
136                crate::bit_util::set_bit(buffer.as_mut(), i * count + j)
137            }
138        }
139        Self {
140            buffer: BooleanBuffer::new(buffer.into(), 0, capacity),
141            null_count: self.null_count * count,
142        }
143    }
144
145    /// Returns the length of this [`NullBuffer`] in bits
146    #[inline]
147    pub fn len(&self) -> usize {
148        self.buffer.len()
149    }
150
151    /// Returns the offset of this [`NullBuffer`] in bits
152    #[inline]
153    pub fn offset(&self) -> usize {
154        self.buffer.offset()
155    }
156
157    /// Returns true if this [`NullBuffer`] is empty
158    #[inline]
159    pub fn is_empty(&self) -> bool {
160        self.buffer.is_empty()
161    }
162
163    /// Free up unused memory.
164    pub fn shrink_to_fit(&mut self) {
165        self.buffer.shrink_to_fit();
166    }
167
168    /// Returns the null count for this [`NullBuffer`]
169    #[inline]
170    pub fn null_count(&self) -> usize {
171        self.null_count
172    }
173
174    /// Returns `true` if the value at `idx` is not null
175    ///
176    /// # Panics
177    ///
178    /// Panics if `idx >= self.len()`
179    #[inline]
180    pub fn is_valid(&self, idx: usize) -> bool {
181        self.buffer.value(idx)
182    }
183
184    /// Returns `true` if the value at `idx` is null
185    ///
186    /// # Panics
187    ///
188    /// Panics if `idx >= self.len()`
189    #[inline]
190    pub fn is_null(&self, idx: usize) -> bool {
191        !self.is_valid(idx)
192    }
193
194    /// Returns the packed validity of this [`NullBuffer`] not including any offset
195    #[inline]
196    pub fn validity(&self) -> &[u8] {
197        self.buffer.values()
198    }
199
200    /// Slices this [`NullBuffer`] by the provided `offset` and `length`
201    ///
202    /// # Panics
203    ///
204    /// Panics if `offset + len > self.len()`
205    pub fn slice(&self, offset: usize, len: usize) -> Self {
206        Self::new(self.buffer.slice(offset, len))
207    }
208
209    /// Returns an iterator over the bits in this [`NullBuffer`]
210    ///
211    /// * `true` indicates that the corresponding value is not NULL
212    /// * `false` indicates that the corresponding value is NULL
213    ///
214    /// Note: [`Self::valid_indices`] will be significantly faster for most use-cases
215    pub fn iter(&self) -> BitIterator<'_> {
216        self.buffer.iter()
217    }
218
219    /// Returns a [`BitIndexIterator`] over the valid indices in this [`NullBuffer`]
220    ///
221    /// Valid indices indicate the corresponding value is not NULL
222    pub fn valid_indices(&self) -> BitIndexIterator<'_> {
223        self.buffer.set_indices()
224    }
225
226    /// Returns a [`BitSliceIterator`] yielding contiguous ranges of valid indices
227    ///
228    /// Valid indices indicate the corresponding value is not NULL
229    pub fn valid_slices(&self) -> BitSliceIterator<'_> {
230        self.buffer.set_slices()
231    }
232
233    /// Calls the provided closure for each index in this null mask that is set
234    #[inline]
235    pub fn try_for_each_valid_idx<E, F: FnMut(usize) -> Result<(), E>>(
236        &self,
237        f: F,
238    ) -> Result<(), E> {
239        if self.null_count == self.len() {
240            return Ok(());
241        }
242        self.valid_indices().try_for_each(f)
243    }
244
245    /// Returns the inner [`BooleanBuffer`]
246    #[inline]
247    pub fn inner(&self) -> &BooleanBuffer {
248        &self.buffer
249    }
250
251    /// Returns the inner [`BooleanBuffer`]
252    #[inline]
253    pub fn into_inner(self) -> BooleanBuffer {
254        self.buffer
255    }
256
257    /// Returns the underlying [`Buffer`]
258    #[inline]
259    pub fn buffer(&self) -> &Buffer {
260        self.buffer.inner()
261    }
262
263    /// Create a [`NullBuffer`] from an *unsliced* validity bitmap (`offset = 0` **bits**) of length `len`.
264    ///
265    /// Returns `None` if there are no nulls (all values valid).
266    pub fn from_unsliced_buffer(buffer: impl Into<Buffer>, len: usize) -> Option<Self> {
267        let bb = BooleanBuffer::new(buffer.into(), 0, len);
268        let nb = NullBuffer::new(bb);
269        (nb.null_count() > 0).then_some(nb)
270    }
271
272    /// Claim memory used by this null buffer in the provided memory pool.
273    #[cfg(feature = "pool")]
274    pub fn claim(&self, pool: &dyn crate::MemoryPool) {
275        // NullBuffer wraps a BooleanBuffer which wraps a Buffer
276        self.buffer.inner().claim(pool);
277    }
278}
279
280impl<'a> IntoIterator for &'a NullBuffer {
281    type Item = bool;
282    type IntoIter = BitIterator<'a>;
283
284    fn into_iter(self) -> Self::IntoIter {
285        self.buffer.iter()
286    }
287}
288
289impl From<BooleanBuffer> for NullBuffer {
290    fn from(value: BooleanBuffer) -> Self {
291        Self::new(value)
292    }
293}
294
295impl From<&[bool]> for NullBuffer {
296    fn from(value: &[bool]) -> Self {
297        BooleanBuffer::from(value).into()
298    }
299}
300
301impl<const N: usize> From<&[bool; N]> for NullBuffer {
302    fn from(value: &[bool; N]) -> Self {
303        value[..].into()
304    }
305}
306
307impl From<Vec<bool>> for NullBuffer {
308    fn from(value: Vec<bool>) -> Self {
309        BooleanBuffer::from(value).into()
310    }
311}
312
313impl FromIterator<bool> for NullBuffer {
314    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {
315        BooleanBuffer::from_iter(iter).into()
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn test_size() {
325        // This tests that the niche optimisation eliminates the overhead of an option
326        assert_eq!(
327            std::mem::size_of::<NullBuffer>(),
328            std::mem::size_of::<Option<NullBuffer>>()
329        );
330    }
331
332    #[test]
333    fn test_from_unsliced_buffer_with_nulls() {
334        // 0b10110010 → null(0), valid(1), null(2), null(3), valid(4), valid(5), null(6), valid(7)
335        let buf = Buffer::from([0b10110010u8]);
336        let result = NullBuffer::from_unsliced_buffer(buf, 8);
337        assert!(result.is_some());
338        let nb = result.unwrap();
339        assert_eq!(nb.len(), 8);
340        assert_eq!(nb.null_count(), 4);
341        assert!(nb.is_null(0));
342        assert!(nb.is_valid(1));
343        assert!(nb.is_null(2));
344        assert!(nb.is_null(3));
345        assert!(nb.is_valid(4));
346        assert!(nb.is_valid(5));
347        assert!(nb.is_null(6));
348        assert!(nb.is_valid(7));
349    }
350
351    #[test]
352    fn test_from_unsliced_buffer_all_valid() {
353        // All bits set = all valid, no nulls
354        let buf = Buffer::from([0b11111111u8]);
355        let result = NullBuffer::from_unsliced_buffer(buf, 8);
356        assert!(result.is_none());
357    }
358
359    #[test]
360    fn test_from_unsliced_buffer_all_null() {
361        // No bits set = all null
362        let buf = Buffer::from([0b00000000u8]);
363        let result = NullBuffer::from_unsliced_buffer(buf, 8);
364        assert!(result.is_some());
365        let nb = result.unwrap();
366        assert_eq!(nb.len(), 8);
367        assert_eq!(nb.null_count(), 8);
368    }
369
370    #[test]
371    fn test_from_unsliced_buffer_empty() {
372        let buf = Buffer::from([]);
373        let result = NullBuffer::from_unsliced_buffer(buf, 0);
374        assert!(result.is_none());
375    }
376
377    #[test]
378    fn test_union_many_all_none() {
379        let result = NullBuffer::union_many([None, None, None]);
380        assert!(result.is_none());
381    }
382
383    #[test]
384    fn test_union_many_single_some() {
385        let a = NullBuffer::from(&[true, false, true, true]);
386        let result = NullBuffer::union_many([Some(&a)]);
387        assert_eq!(result, Some(a));
388    }
389
390    #[test]
391    fn test_union_many_two_inputs() {
392        let a = NullBuffer::from(&[true, false, true, true]);
393        let b = NullBuffer::from(&[true, true, false, true]);
394        let result = NullBuffer::union_many([Some(&a), Some(&b)]);
395        let expected = NullBuffer::union(Some(&a), Some(&b));
396        assert_eq!(result, expected);
397    }
398
399    #[test]
400    fn test_union_many_three_inputs() {
401        let a = NullBuffer::from(&[true, false, true, true]);
402        let b = NullBuffer::from(&[true, true, false, true]);
403        let c = NullBuffer::from(&[false, true, true, true]);
404        let result = NullBuffer::union_many([Some(&a), Some(&b), Some(&c)]);
405        let expected = NullBuffer::from(&[false, false, false, true]);
406        assert_eq!(result, Some(expected));
407    }
408
409    #[test]
410    fn test_union_many_mixed_none() {
411        let a = NullBuffer::from(&[true, false, true, true]);
412        let b = NullBuffer::from(&[false, true, true, true]);
413        let result = NullBuffer::union_many([Some(&a), None, Some(&b)]);
414        let expected = NullBuffer::union(Some(&a), Some(&b));
415        assert_eq!(result, expected);
416    }
417
418    #[test]
419    fn test_union_many_empty_slice() {
420        let result = NullBuffer::union_many([] as [Option<&NullBuffer>; 0]);
421        assert!(result.is_none());
422    }
423
424    #[test]
425    fn test_union_many_no_nulls() {
426        let a = NullBuffer::from(&[true, true, true, true]);
427
428        let result = NullBuffer::union_many([Some(&a), Some(&a), Some(&a)]);
429        assert_eq!(result, None);
430    }
431
432    #[test]
433    fn test_union_no_nulls() {
434        let a = NullBuffer::from(&[true, true, true, true]);
435
436        let result = NullBuffer::union(Some(&a), Some(&a));
437        assert_eq!(result, None);
438
439        let result = NullBuffer::union(Some(&a), None);
440        assert_eq!(result, None);
441
442        let result = NullBuffer::union(None, Some(&a));
443        assert_eq!(result, None);
444    }
445
446    #[test]
447    fn test_union_nulls_one_side() {
448        let all_valid = NullBuffer::from(&[true, true, true, true]);
449        let all_null = NullBuffer::from(&[false, false, false, false]);
450
451        let result = NullBuffer::union(Some(&all_valid), Some(&all_null));
452        assert_eq!(result, Some(all_null.clone()));
453
454        let result = NullBuffer::union(Some(&all_null), Some(&all_valid));
455        assert_eq!(result, Some(all_null.clone()));
456    }
457}
arrow_buffer/buffer/null.rs

arrow_buffer/buffer/
null.rs
