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::ToCanonical;
11use crate::arrays::BoolArray;
12use crate::arrays::BoolVTable;
13use crate::expr::FillNullKernel;
14use crate::scalar::Scalar;
15use crate::validity::Validity;
16use crate::vtable::ValidityHelper;
17
18impl FillNullKernel for BoolVTable {
19    fn fill_null(
20        array: &BoolArray,
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 bool_buffer = if fill {
32                    array.to_bit_buffer() | &!v.to_bool().to_bit_buffer()
33                } else {
34                    array.to_bit_buffer() & v.to_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    use vortex_dtype::DType;
49    use vortex_dtype::Nullability;
50
51    use crate::arrays::BoolArray;
52    use crate::builtins::ArrayBuiltins;
53    use crate::canonical::ToCanonical;
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
