Skip to main content

arrow_digest/
utils.rs

1use arrow::buffer::{BooleanBuffer, NullBuffer, buffer_bin_or};
2
3/////////////////////////////////////////////////////////////////////////////////////////
4
5pub(crate) fn combine_null_buffers(a: &NullBuffer, b: &NullBuffer) -> NullBuffer {
6    assert_eq!(
7        a.len(),
8        b.len(),
9        "Attempting to combine buffers of different size {} != {}",
10        a.len(),
11        b.len()
12    );
13
14    // In NullBuffer `1` stands for valid and `0` for null
15    let buffer = buffer_bin_or(a.buffer(), a.offset(), b.buffer(), b.offset(), a.len());
16
17    NullBuffer::new(BooleanBuffer::new(buffer, 0, a.len()))
18}
19
20/////////////////////////////////////////////////////////////////////////////////////////
21
22pub(crate) fn maybe_combine_null_buffers<'a>(
23    a: Option<&'a NullBuffer>,
24    b: Option<&'a NullBuffer>,
25) -> CombinedNullBuffer<'a> {
26    match (a, b) {
27        (None, Some(b)) if b.null_count() != 0 => CombinedNullBuffer::Borrowed(b),
28        (Some(a), None) if a.null_count() != 0 => CombinedNullBuffer::Borrowed(a),
29        (Some(a), Some(b)) if a.null_count() != 0 && b.null_count() != 0 => {
30            CombinedNullBuffer::Owned(combine_null_buffers(a, b))
31        }
32        _ => CombinedNullBuffer::None,
33    }
34}
35
36/////////////////////////////////////////////////////////////////////////////////////////
37
38// An amalgamation of Option<T> and a Cow<'a, &T>
39pub(crate) enum CombinedNullBuffer<'a> {
40    None,
41    Borrowed(&'a NullBuffer),
42    Owned(NullBuffer),
43}
44
45impl<'a> CombinedNullBuffer<'a> {
46    pub fn as_option(&'a self) -> Option<&'a NullBuffer> {
47        match self {
48            CombinedNullBuffer::None => None,
49            CombinedNullBuffer::Borrowed(buf) => Some(*buf),
50            CombinedNullBuffer::Owned(buf) => Some(buf),
51        }
52    }
53}