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

vortex_array/arrays/bool/compute/
fill_null.rs
