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::expr::FillNullReduce;
12use crate::scalar::Scalar;
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    use vortex_dtype::DType;
34    use vortex_dtype::Nullability;
35
36    use crate::array::Array;
37    use crate::arrays::BoolArray;
38    use crate::arrays::ChunkedArray;
39    use crate::builtins::ArrayBuiltins;
40    use crate::scalar::Scalar;
41    use crate::validity::Validity;
42
43    #[test]
44    fn fill_null_chunks() {
45        let chunked = ChunkedArray::try_new(
46            vec![
47                BoolArray::new(BitBuffer::new_set(5), Validity::AllInvalid).to_array(),
48                BoolArray::new(BitBuffer::new_set(5), Validity::AllValid).to_array(),
49            ],
50            DType::Bool(Nullability::Nullable),
51        )
52        .unwrap();
53
54        let filled = chunked.to_array().fill_null(Scalar::from(false)).unwrap();
55        assert_eq!(*filled.dtype(), DType::Bool(Nullability::NonNullable));
56    }
57}