vortex_dict/compute/
fill_null.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::compute::{FillNullFn, Operator, compare, fill_null};
3use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
4use vortex_error::VortexResult;
5use vortex_scalar::{Scalar, ScalarValue};
6
7use crate::{DictArray, DictEncoding};
8
9impl FillNullFn<&DictArray> for DictEncoding {
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()),
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.codes().dtype().clone(),
32                ScalarValue::from(first_fill_value),
33            ),
34        )?;
35        // And fill nulls in the values
36        let values = fill_null(array.values(), fill_value)?;
37
38        Ok(DictArray::try_new(codes, values)?.into_array())
39    }
40}
vortex_dict/compute/fill_null.rs

vortex_dict/compute/
fill_null.rs
