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;
5use vortex_error::vortex_err;
6
7use crate::ArrayRef;
8use crate::ExecutionCtx;
9use crate::IntoArray;
10use crate::array::ArrayView;
11use crate::arrays::Bool;
12use crate::arrays::BoolArray;
13use crate::arrays::bool::BoolArrayExt;
14use crate::scalar::Scalar;
15use crate::scalar_fn::fns::fill_null::FillNullKernel;
16use crate::validity::Validity;
17
18impl FillNullKernel for Bool {
19    fn fill_null(
20        array: ArrayView<'_, Bool>,
21        fill_value: &Scalar,
22        ctx: &mut ExecutionCtx,
23    ) -> VortexResult<Option<ArrayRef>> {
24        let fill = fill_value
25            .as_bool()
26            .value()
27            .ok_or_else(|| vortex_err!("Fill value must be non null"))?;
28
29        Ok(Some(match array.validity()? {
30            Validity::Array(v) => {
31                let v_bool = v.execute::<BoolArray>(ctx)?;
32                let bool_buffer = if fill {
33                    array.to_bit_buffer() | &!v_bool.to_bit_buffer()
34                } else {
35                    array.to_bit_buffer() & v_bool.to_bit_buffer()
36                };
37                BoolArray::new(bool_buffer, fill_value.dtype().nullability().into()).into_array()
38            }
39            _ => unreachable!("checked in entry point"),
40        }))
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use rstest::rstest;
47    use vortex_buffer::BitBuffer;
48    use vortex_buffer::bitbuffer;
49
50    use crate::IntoArray;
51    use crate::arrays::BoolArray;
52    use crate::arrays::bool::BoolArrayExt;
53    use crate::builtins::ArrayBuiltins;
54    use crate::canonical::ToCanonical;
55    use crate::dtype::DType;
56    use crate::dtype::Nullability;
57    use crate::scalar::Scalar;
58    use crate::validity::Validity;
59
60    #[rstest]
61    #[case(true, bitbuffer![true, true, false, true])]
62    #[case(false, bitbuffer![true, false, false, false])]
63    fn bool_fill_null(#[case] fill_value: bool, #[case] expected: BitBuffer) {
64        let bool_array = BoolArray::new(
65            BitBuffer::from_iter([true, true, false, false]),
66            Validity::from_iter([true, false, true, false]),
67        );
68        let non_null_array = bool_array
69            .into_array()
70            .fill_null(Scalar::from(fill_value))
71            .unwrap()
72            .to_bool();
73        assert_eq!(non_null_array.to_bit_buffer(), expected);
74        assert_eq!(
75            non_null_array.dtype(),
76            &DType::Bool(Nullability::NonNullable)
77        );
78    }
79}
vortex_array/arrays/bool/compute/fill_null.rs

vortex_array/arrays/bool/compute/
fill_null.rs
