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

vortex_array/arrays/chunked/compute/
fill_null.rs
