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::from_bool_buffer(
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::from_bool_buffer(bool_buffer, fill_value.dtype().nullability().into())
36                    .into_array()
37            }
38        })
39    }
40}
41
42register_kernel!(FillNullKernelAdapter(BoolVTable).lift());
43
44#[cfg(test)]
45mod tests {
46    use arrow_buffer::BooleanBuffer;
47    use rstest::rstest;
48    use vortex_dtype::{DType, Nullability};
49
50    use crate::arrays::BoolArray;
51    use crate::canonical::ToCanonical;
52    use crate::compute::fill_null;
53    use crate::validity::Validity;
54
55    #[rstest]
56    #[case(true, vec![true, true, false, true])]
57    #[case(false, vec![true, false, false, false])]
58    fn bool_fill_null(#[case] fill_value: bool, #[case] expected: Vec<bool>) {
59        let bool_array = BoolArray::from_bool_buffer(
60            BooleanBuffer::from_iter([true, true, false, false]),
61            Validity::from_iter([true, false, true, false]),
62        );
63        let non_null_array = fill_null(bool_array.as_ref(), &fill_value.into())
64            .unwrap()
65            .to_bool();
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
