vortex_array/arrays/bool/compute/
fill_null.rs

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

vortex_array/arrays/bool/compute/
fill_null.rs
