vortex-array 0.68.0

Vortex in memory columnar data format
Documentation
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use arrow_buffer::BooleanBuffer;
use arrow_buffer::NullBuffer;
use vortex_error::VortexResult;
use vortex_mask::Mask;

use crate::ExecutionCtx;
use crate::validity::Validity;

/// Converts a [`Validity`] to an Arrow [`NullBuffer`], executing the validity array if needed.
pub fn to_arrow_null_buffer(
    validity: Validity,
    len: usize,
    ctx: &mut ExecutionCtx,
) -> VortexResult<Option<NullBuffer>> {
    Ok(match validity {
        Validity::NonNullable | Validity::AllValid => None,
        Validity::AllInvalid => Some(NullBuffer::new_null(len)),
        Validity::Array(array) => to_null_buffer(array.execute::<Mask>(ctx)?),
    })
}

/// Converts a mask to a null buffer.
pub fn to_null_buffer(mask: Mask) -> Option<NullBuffer> {
    match mask {
        Mask::AllTrue(_) => None,
        Mask::AllFalse(l) => Some(NullBuffer::new_null(l)),
        Mask::Values(values) => Some(NullBuffer::from(BooleanBuffer::from(
            values.bit_buffer().clone(),
        ))),
    }
}

#[cfg(test)]
mod tests {
    use vortex_buffer::BitBuffer;
    use vortex_mask::Mask;

    use crate::arrow::null_buffer::to_null_buffer;

    #[test]
    fn test_mask_to_null_buffer() {
        let all_true = Mask::new_true(5);
        assert!(to_null_buffer(all_true).is_none());

        let all_false = Mask::new_false(5);
        let null_buffer = to_null_buffer(all_false).unwrap();
        assert_eq!(null_buffer.null_count(), 5);

        let values = Mask::from_buffer(BitBuffer::from_iter([true, false, true, false, true]));
        let null_buffer = to_null_buffer(values).unwrap();
        assert_eq!(null_buffer.null_count(), 2);
    }
}