1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

use vortex_error::{vortex_err, VortexResult};
use vortex_scalar::Scalar;

use crate::array::{BoolArray, BoolEncoding, ConstantArray};
use crate::compute::FillNullFn;
use crate::validity::Validity;
use crate::{ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant};

impl FillNullFn<BoolArray> for BoolEncoding {
    fn fill_null(&self, array: &BoolArray, fill_value: Scalar) -> VortexResult<ArrayData> {
        let fill = fill_value
            .as_bool()
            .value()
            .ok_or_else(|| vortex_err!("Fill value must be non null"))?;

        Ok(match array.validity() {
            Validity::NonNullable => array.clone().into_array(),
            Validity::AllValid => BoolArray::from(array.boolean_buffer()).into_array(),
            Validity::AllInvalid => ConstantArray::new(fill, array.len()).into_array(),
            Validity::Array(v) => {
                let bool_buffer = if fill {
                    &array.boolean_buffer() | &!&v.into_bool()?.boolean_buffer()
                } else {
                    &array.boolean_buffer() & &v.into_bool()?.boolean_buffer()
                };
                BoolArray::from(bool_buffer).into_array()
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use arrow_buffer::BooleanBuffer;
    use rstest::rstest;
    use vortex_dtype::{DType, Nullability};

    use crate::array::BoolArray;
    use crate::compute::fill_null;
    use crate::validity::Validity;
    use crate::{ArrayDType, IntoArrayVariant};

    #[rstest]
    #[case(true, vec![true, true, false, true])]
    #[case(false, vec![true, false, false, false])]
    fn bool_fill_null(#[case] fill_value: bool, #[case] expected: Vec<bool>) {
        let bool_array = BoolArray::try_new(
            BooleanBuffer::from_iter([true, true, false, false]),
            Validity::from_iter([true, false, true, false]),
        )
        .unwrap();
        let non_null_array = fill_null(bool_array, fill_value.into())
            .unwrap()
            .into_bool()
            .unwrap();
        assert_eq!(
            non_null_array.boolean_buffer().iter().collect::<Vec<_>>(),
            expected
        );
        assert_eq!(
            non_null_array.dtype(),
            &DType::Bool(Nullability::NonNullable)
        );
    }
}