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::VortexSessionExecute;
40    use crate::array_session;
41    use crate::arrays::BoolArray;
42    use crate::arrays::ChunkedArray;
43    use crate::assert_arrays_eq;
44    use crate::builtins::ArrayBuiltins;
45    use crate::dtype::DType;
46    use crate::dtype::Nullability;
47    use crate::scalar::Scalar;
48    use crate::validity::Validity;
49
50    #[test]
51    fn fill_null_chunks() {
52        let chunked = ChunkedArray::try_new(
53            vec![
54                BoolArray::new(BitBuffer::new_set(5), Validity::AllInvalid).into_array(),
55                BoolArray::new(BitBuffer::new_set(5), Validity::AllValid).into_array(),
56            ],
57            DType::Bool(Nullability::Nullable),
58        )
59        .unwrap();
60
61        let filled = chunked.into_array().fill_null(Scalar::from(false)).unwrap();
62        assert_arrays_eq!(
63            filled,
64            BoolArray::from_iter([
65                false, false, false, false, false, true, true, true, true, true
66            ]),
67            &mut array_session().create_execution_ctx()
68        );
69    }
70}
vortex_array/arrays/chunked/compute/fill_null.rs

vortex_array/arrays/chunked/compute/
fill_null.rs
