vortex_array/arrays/chunked/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_scalar::Scalar;
6
7use crate::arrays::{ChunkedArray, ChunkedVTable};
8use crate::compute::{FillNullKernel, FillNullKernelAdapter, fill_null};
9use crate::{ArrayRef, IntoArray, register_kernel};
10
11impl FillNullKernel for ChunkedVTable {
12    fn fill_null(&self, array: &ChunkedArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
13        Ok(ChunkedArray::new_unchecked(
14            array
15                .chunks()
16                .iter()
17                .map(|c| fill_null(c, fill_value))
18                .collect::<VortexResult<Vec<_>>>()?,
19            fill_value.dtype().clone(),
20        )
21        .into_array())
22    }
23}
24
25register_kernel!(FillNullKernelAdapter(ChunkedVTable).lift());
26
27#[cfg(test)]
28mod tests {
29    use arrow_buffer::BooleanBuffer;
30    use vortex_dtype::{DType, Nullability};
31
32    use crate::array::Array;
33    use crate::arrays::{BoolArray, ChunkedArray};
34    use crate::compute::fill_null;
35    use crate::validity::Validity;
36
37    #[test]
38    fn fill_null_chunks() {
39        let chunked = ChunkedArray::try_new(
40            vec![
41                BoolArray::new(BooleanBuffer::new_set(5), Validity::AllInvalid).to_array(),
42                BoolArray::new(BooleanBuffer::new_set(5), Validity::AllValid).to_array(),
43            ],
44            DType::Bool(Nullability::Nullable),
45        )
46        .unwrap();
47
48        let filled = fill_null(chunked.as_ref(), &false.into()).unwrap();
49        assert_eq!(*filled.dtype(), DType::Bool(Nullability::NonNullable));
50    }
51}