vortex_array/arrays/bool/compute/
fill_null.rs

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

vortex_array/arrays/bool/compute/
fill_null.rs
