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