vortex_array/arrays/bool/compute/
fill_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::{VortexResult, vortex_err};
5use vortex_scalar::Scalar;
6
7use crate::arrays::{BoolArray, BoolVTable};
8use crate::compute::{FillNullKernel, FillNullKernelAdapter};
9use crate::validity::Validity;
10use crate::vtable::ValidityHelper;
11use crate::{ArrayRef, IntoArray, ToCanonical, register_kernel};
12
13impl FillNullKernel for BoolVTable {
14    fn fill_null(&self, array: &BoolArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
15        let fill = fill_value
16            .as_bool()
17            .value()
18            .ok_or_else(|| vortex_err!("Fill value must be non null"))?;
19
20        Ok(match array.validity() {
21            Validity::Array(v) => {
22                let bool_buffer = if fill {
23                    array.bit_buffer() | &!v.to_bool().bit_buffer()
24                } else {
25                    array.bit_buffer() & v.to_bool().bit_buffer()
26                };
27                BoolArray::from_bit_buffer(bool_buffer, fill_value.dtype().nullability().into())
28                    .into_array()
29            }
30            _ => unreachable!("checked in entry point"),
31        })
32    }
33}
34
35register_kernel!(FillNullKernelAdapter(BoolVTable).lift());
36
37#[cfg(test)]
38mod tests {
39    use rstest::rstest;
40    use vortex_buffer::{BitBuffer, bitbuffer};
41    use vortex_dtype::{DType, Nullability};
42
43    use crate::arrays::BoolArray;
44    use crate::canonical::ToCanonical;
45    use crate::compute::fill_null;
46    use crate::validity::Validity;
47
48    #[rstest]
49    #[case(true, bitbuffer![true, true, false, true])]
50    #[case(false, bitbuffer![true, false, false, false])]
51    fn bool_fill_null(#[case] fill_value: bool, #[case] expected: BitBuffer) {
52        let bool_array = BoolArray::from_bit_buffer(
53            BitBuffer::from_iter([true, true, false, false]),
54            Validity::from_iter([true, false, true, false]),
55        );
56        let non_null_array = fill_null(bool_array.as_ref(), &fill_value.into())
57            .unwrap()
58            .to_bool();
59        assert_eq!(non_null_array.bit_buffer(), &expected);
60        assert_eq!(
61            non_null_array.dtype(),
62            &DType::Bool(Nullability::NonNullable)
63        );
64    }
65}
vortex_array/arrays/bool/compute/fill_null.rs

vortex_array/arrays/bool/compute/
fill_null.rs
