Skip to main content

vortex_array/arrow/
null_buffer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_buffer::BooleanBuffer;
5use arrow_buffer::NullBuffer;
6use vortex_error::VortexResult;
7use vortex_mask::Mask;
8
9use crate::ExecutionCtx;
10use crate::validity::Validity;
11
12/// Converts a [`Validity`] to an Arrow [`NullBuffer`], executing the validity array if needed.
13pub fn to_arrow_null_buffer(
14    validity: Validity,
15    len: usize,
16    ctx: &mut ExecutionCtx,
17) -> VortexResult<Option<NullBuffer>> {
18    Ok(match validity {
19        Validity::NonNullable | Validity::AllValid => None,
20        Validity::AllInvalid => Some(NullBuffer::new_null(len)),
21        Validity::Array(array) => to_null_buffer(array.execute::<Mask>(ctx)?),
22    })
23}
24
25/// Converts a mask to a null buffer.
26pub fn to_null_buffer(mask: Mask) -> Option<NullBuffer> {
27    match mask {
28        Mask::AllTrue(_) => None,
29        Mask::AllFalse(l) => Some(NullBuffer::new_null(l)),
30        Mask::Values(values) => Some(NullBuffer::from(BooleanBuffer::from(
31            values.bit_buffer().clone(),
32        ))),
33    }
34}
35
36#[cfg(test)]
37mod tests {
38    use vortex_buffer::BitBuffer;
39    use vortex_mask::Mask;
40
41    use crate::arrow::null_buffer::to_null_buffer;
42
43    #[test]
44    fn test_mask_to_null_buffer() {
45        let all_true = Mask::new_true(5);
46        assert!(to_null_buffer(all_true).is_none());
47
48        let all_false = Mask::new_false(5);
49        let null_buffer = to_null_buffer(all_false).unwrap();
50        assert_eq!(null_buffer.null_count(), 5);
51
52        let values = Mask::from_buffer(BitBuffer::from_iter([true, false, true, false, true]));
53        let null_buffer = to_null_buffer(values).unwrap();
54        assert_eq!(null_buffer.null_count(), 2);
55    }
56}