vortex_array/compute/
fill_null.rs

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

vortex_array/compute/
fill_null.rs
