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

vortex_array/arrays/bool/compute/
fill_null.rs
