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, ConstantArray};
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::NonNullable | Validity::AllValid => BoolArray::new(
22                array.boolean_buffer().clone(),
23                fill_value.dtype().nullability().into(),
24            )
25            .into_array(),
26            Validity::AllInvalid => {
27                ConstantArray::new(fill_value.clone(), array.len()).into_array()
28            }
29            Validity::Array(v) => {
30                let bool_buffer = if fill {
31                    array.boolean_buffer() | &!v.to_bool()?.boolean_buffer()
32                } else {
33                    array.boolean_buffer() & v.to_bool()?.boolean_buffer()
34                };
35                BoolArray::new(bool_buffer, fill_value.dtype().nullability().into()).into_array()
36            }
37        })
38    }
39}
40
41register_kernel!(FillNullKernelAdapter(BoolVTable).lift());
42
43#[cfg(test)]
44mod tests {
45    use arrow_buffer::BooleanBuffer;
46    use rstest::rstest;
47    use vortex_dtype::{DType, Nullability};
48
49    use crate::arrays::BoolArray;
50    use crate::canonical::ToCanonical;
51    use crate::compute::fill_null;
52    use crate::validity::Validity;
53
54    #[rstest]
55    #[case(true, vec![true, true, false, true])]
56    #[case(false, vec![true, false, false, false])]
57    fn bool_fill_null(#[case] fill_value: bool, #[case] expected: Vec<bool>) {
58        let bool_array = BoolArray::new(
59            BooleanBuffer::from_iter([true, true, false, false]),
60            Validity::from_iter([true, false, true, false]),
61        );
62        let non_null_array = fill_null(bool_array.as_ref(), &fill_value.into())
63            .unwrap()
64            .to_bool()
65            .unwrap();
66        assert_eq!(
67            non_null_array.boolean_buffer().iter().collect::<Vec<_>>(),
68            expected
69        );
70        assert_eq!(
71            non_null_array.dtype(),
72            &DType::Bool(Nullability::NonNullable)
73        );
74    }
75}
vortex_array/arrays/bool/compute/fill_null.rs

vortex_array/arrays/bool/compute/
fill_null.rs
