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