vortex_dict/compute/
fill_null.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::compute::{FillNullKernel, FillNullKernelAdapter, Operator, compare, fill_null};
3use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
4use vortex_error::VortexResult;
5use vortex_scalar::{Scalar, ScalarValue};
6
7use crate::{DictArray, DictVTable};
8
9impl FillNullKernel for DictVTable {
10    fn fill_null(&self, array: &DictArray, fill_value: &Scalar) -> VortexResult<ArrayRef> {
11        // If the fill value exists in the dictionary, we can simply rewrite the null codes to
12        // point to the value.
13        let found_fill_values = compare(
14            array.values(),
15            ConstantArray::new(fill_value.clone(), array.values().len()).as_ref(),
16            Operator::Eq,
17        )?
18        .to_bool()?;
19
20        let Some(first_fill_value) = found_fill_values.boolean_buffer().set_indices().next() else {
21            // No fill values found, so we must canonicalize and fill_null.
22            // TODO(ngates): compute kernels should all return Option<ArrayRef> to support this
23            //  fall back.
24            return fill_null(&array.to_canonical()?.into_array(), fill_value);
25        };
26
27        // Now we rewrite the nullable codes to point at the fill value.
28        let codes = fill_null(
29            array.codes(),
30            &Scalar::new(
31                array
32                    .codes()
33                    .dtype()
34                    .with_nullability(fill_value.dtype().nullability()),
35                ScalarValue::from(first_fill_value),
36            ),
37        )?;
38        // And fill nulls in the values
39        let values = fill_null(array.values(), fill_value)?;
40
41        Ok(DictArray::try_new(codes, values)?.into_array())
42    }
43}
44
45register_kernel!(FillNullKernelAdapter(DictVTable).lift());
46
47#[cfg(test)]
48mod tests {
49    use arrow_buffer::BooleanBuffer;
50    use vortex_array::arrays::PrimitiveArray;
51    use vortex_array::compute::fill_null;
52    use vortex_array::validity::Validity;
53    use vortex_array::{IntoArray, ToCanonical};
54    use vortex_buffer::buffer;
55    use vortex_dtype::Nullability;
56    use vortex_error::VortexUnwrap;
57    use vortex_scalar::Scalar;
58
59    use crate::DictArray;
60
61    #[test]
62    fn nullable_codes_fill_in_values() {
63        let dict = DictArray::try_new(
64            PrimitiveArray::new(
65                buffer![0u32, 1, 2],
66                Validity::from(BooleanBuffer::from(vec![true, false, true])),
67            )
68            .into_array(),
69            PrimitiveArray::new(buffer![10, 20, 20], Validity::AllValid).into_array(),
70        )
71        .vortex_unwrap();
72
73        let filled = fill_null(
74            dict.as_ref(),
75            &Scalar::primitive(20, Nullability::NonNullable),
76        )
77        .vortex_unwrap();
78        let filled_primitive = filled.to_primitive().vortex_unwrap();
79        assert_eq!(filled_primitive.as_slice::<i32>(), [10, 20, 20]);
80        assert!(filled_primitive.all_valid().vortex_unwrap());
81    }
82}
vortex_dict/compute/fill_null.rs

vortex_dict/compute/
fill_null.rs
