vortex_array/arrays/chunked/compute/
fill_null.rs

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