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

vortex_array/arrays/chunked/compute/
fill_null.rs
