vortex_array/arrays/chunked/compute/
fill_null.rs

1use vortex_error::VortexResult;
2use vortex_scalar::Scalar;
3
4use crate::arrays::{ChunkedArray, ChunkedEncoding};
5use crate::compute::{FillNullFn, fill_null};
6use crate::{Array, ArrayRef};
7
8impl FillNullFn<&ChunkedArray> for ChunkedEncoding {
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.clone()))
15                .collect::<VortexResult<Vec<_>>>()?,
16            array.dtype().as_nonnullable(),
17        )
18        .into_array())
19    }
20}
21
22#[cfg(test)]
23mod tests {
24    use arrow_buffer::BooleanBuffer;
25    use vortex_dtype::{DType, Nullability};
26
27    use crate::array::Array;
28    use crate::arrays::{BoolArray, ChunkedArray};
29    use crate::compute::fill_null;
30    use crate::validity::Validity;
31
32    #[test]
33    fn fill_null_chunks() {
34        let chunked = ChunkedArray::try_new(
35            vec![
36                BoolArray::new(BooleanBuffer::new_set(5), Validity::AllInvalid).to_array(),
37                BoolArray::new(BooleanBuffer::new_set(5), Validity::AllValid).to_array(),
38            ],
39            DType::Bool(Nullability::Nullable),
40        )
41        .unwrap();
42
43        let filled = fill_null(&chunked, false.into()).unwrap();
44        assert_eq!(*filled.dtype(), DType::Bool(Nullability::NonNullable));
45    }
46}