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 => array.to_array().into_array(),
18            Validity::AllValid => BoolArray::from(array.boolean_buffer().clone()).into_array(),
19            Validity::AllInvalid => ConstantArray::new(fill, array.len()).into_array(),
20            Validity::Array(v) => {
21                let bool_buffer = if fill {
22                    array.boolean_buffer() | &!v.to_bool()?.boolean_buffer()
23                } else {
24                    array.boolean_buffer() & v.to_bool()?.boolean_buffer()
25                };
26                BoolArray::from(bool_buffer).into_array()
27            }
28        })
29    }
30}
31
32#[cfg(test)]
33mod tests {
34    use arrow_buffer::BooleanBuffer;
35    use rstest::rstest;
36    use vortex_dtype::{DType, Nullability};
37
38    use crate::array::Array;
39    use crate::arrays::BoolArray;
40    use crate::canonical::ToCanonical;
41    use crate::compute::fill_null;
42    use crate::validity::Validity;
43
44    #[rstest]
45    #[case(true, vec![true, true, false, true])]
46    #[case(false, vec![true, false, false, false])]
47    fn bool_fill_null(#[case] fill_value: bool, #[case] expected: Vec<bool>) {
48        let bool_array = BoolArray::new(
49            BooleanBuffer::from_iter([true, true, false, false]),
50            Validity::from_iter([true, false, true, false]),
51        );
52        let non_null_array = fill_null(&bool_array, fill_value.into())
53            .unwrap()
54            .to_bool()
55            .unwrap();
56        assert_eq!(
57            non_null_array.boolean_buffer().iter().collect::<Vec<_>>(),
58            expected
59        );
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
