Skip to main content

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