vortex_array/compute/
fill_null.rs

1use vortex_error::{VortexExpect, VortexResult, vortex_bail};
2use vortex_scalar::Scalar;
3
4use crate::compute::try_cast;
5use crate::encoding::Encoding;
6use crate::{Array, ArrayRef, IntoArray};
7
8/// Implementation of fill_null for an encoding.
9///
10/// SAFETY: the fill value is guaranteed to be non-null.
11pub trait FillNullFn<A> {
12    fn fill_null(&self, array: A, fill_value: Scalar) -> VortexResult<ArrayRef>;
13}
14
15impl<E: Encoding> FillNullFn<&dyn Array> for E
16where
17    E: for<'a> FillNullFn<&'a E::Array>,
18{
19    fn fill_null(&self, array: &dyn Array, fill_value: Scalar) -> VortexResult<ArrayRef> {
20        let array_ref = array
21            .as_any()
22            .downcast_ref::<E::Array>()
23            .vortex_expect("Failed to downcast array");
24        FillNullFn::fill_null(self, array_ref, fill_value)
25    }
26}
27
28pub fn fill_null(array: &dyn Array, fill_value: Scalar) -> VortexResult<ArrayRef> {
29    if !array.dtype().is_nullable() {
30        return Ok(array.to_array());
31    }
32
33    if array.invalid_count()? == 0 {
34        return try_cast(array, fill_value.dtype());
35    }
36
37    if fill_value.is_null() {
38        vortex_bail!("Cannot fill_null with a null value")
39    }
40
41    if !array.dtype().eq_ignore_nullability(fill_value.dtype()) {
42        vortex_bail!(MismatchedTypes: array.dtype(), fill_value.dtype())
43    }
44
45    let fill_value_nullability = fill_value.dtype().nullability();
46    let filled = fill_null_impl(array, fill_value)?;
47
48    assert_eq!(
49        filled.len(),
50        array.len(),
51        "FillNull length mismatch {}",
52        array.encoding()
53    );
54    assert_eq!(
55        filled.dtype().nullability(),
56        fill_value_nullability,
57        "FillNull nullability mismatch {}",
58        array.encoding()
59    );
60
61    Ok(filled)
62}
63
64fn fill_null_impl(array: &dyn Array, fill_value: Scalar) -> VortexResult<ArrayRef> {
65    if let Some(fill_null_fn) = array.vtable().fill_null_fn() {
66        return fill_null_fn.fill_null(array, fill_value);
67    }
68
69    log::debug!("FillNullFn not implemented for {}", array.encoding());
70    let canonical_arr = array.to_canonical()?.into_array();
71    if let Some(fill_null_fn) = canonical_arr.vtable().fill_null_fn() {
72        return fill_null_fn.fill_null(&canonical_arr, fill_value);
73    }
74
75    vortex_bail!(
76        "fill null not implemented for canonical encoding {}, fallback from {}",
77        canonical_arr.encoding(),
78        array.encoding()
79    )
80}
vortex_array/compute/fill_null.rs

vortex_array/compute/
fill_null.rs
