Skip to main content

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;
5
6use crate::ArrayRef;
7use crate::IntoArray;
8use crate::arrays::ChunkedArray;
9use crate::arrays::ChunkedVTable;
10use crate::builtins::ArrayBuiltins;
11use crate::scalar::Scalar;
12use crate::scalar_fn::fns::fill_null::FillNullReduce;
13
14impl FillNullReduce for ChunkedVTable {
15    fn fill_null(array: &ChunkedArray, fill_value: &Scalar) -> VortexResult<Option<ArrayRef>> {
16        let new_chunks = array
17            .chunks()
18            .iter()
19            .map(|c| c.fill_null(fill_value.clone()))
20            .collect::<VortexResult<Vec<_>>>()?;
21
22        // SAFETY: wrapping each chunk in ScalarFnArray preserves the same DType across all chunks.
23        Ok(Some(
24            unsafe { ChunkedArray::new_unchecked(new_chunks, fill_value.dtype().clone()) }
25                .into_array(),
26        ))
27    }
28}
29
30#[cfg(test)]
31mod tests {
32    use vortex_buffer::BitBuffer;
33
34    use crate::IntoArray;
35    use crate::array::DynArray;
36    use crate::arrays::BoolArray;
37    use crate::arrays::ChunkedArray;
38    use crate::builtins::ArrayBuiltins;
39    use crate::dtype::DType;
40    use crate::dtype::Nullability;
41    use crate::scalar::Scalar;
42    use crate::validity::Validity;
43
44    #[test]
45    fn fill_null_chunks() {
46        let chunked = ChunkedArray::try_new(
47            vec![
48                BoolArray::new(BitBuffer::new_set(5), Validity::AllInvalid).into_array(),
49                BoolArray::new(BitBuffer::new_set(5), Validity::AllValid).into_array(),
50            ],
51            DType::Bool(Nullability::Nullable),
52        )
53        .unwrap();
54
55        let filled = chunked.into_array().fill_null(Scalar::from(false)).unwrap();
56        assert_eq!(*filled.dtype(), DType::Bool(Nullability::NonNullable));
57    }
58}