1use arrow::buffer::{BooleanBuffer, NullBuffer, buffer_bin_or};
2
3pub(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 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
20pub(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
36pub(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}