vortex_array/compute/
fill_null.rs

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

vortex_array/compute/
fill_null.rs
