vortex_array/compute/
fill_forward.rs

1use vortex_error::{VortexExpect, VortexResult, vortex_err};
2
3use crate::encoding::Encoding;
4use crate::{Array, ArrayRef};
5
6/// Trait for filling forward on an array, i.e., replacing nulls with the last non-null value.
7///
8/// If the array is non-nullable, it is returned as-is.
9/// If the array is entirely nulls, the fill forward operation returns an array of the same length, filled with the default value of the array's type.
10/// The DType of the returned array is the same as the input array; the Validity of the returned array is always either NonNullable or AllValid.
11pub trait FillForwardFn<A> {
12    fn fill_forward(&self, array: A) -> VortexResult<ArrayRef>;
13}
14
15impl<E: Encoding> FillForwardFn<&dyn Array> for E
16where
17    E: for<'a> FillForwardFn<&'a E::Array>,
18{
19    fn fill_forward(&self, array: &dyn Array) -> VortexResult<ArrayRef> {
20        let array_ref = array
21            .as_any()
22            .downcast_ref::<E::Array>()
23            .vortex_expect("Failed to downcast array");
24        FillForwardFn::fill_forward(self, array_ref)
25    }
26}
27
28pub fn fill_forward(array: &dyn Array) -> VortexResult<ArrayRef> {
29    if !array.dtype().is_nullable() {
30        return Ok(array.to_array());
31    }
32
33    let filled = array
34        .vtable()
35        .fill_forward_fn()
36        .map(|f| f.fill_forward(array))
37        .unwrap_or_else(|| {
38            Err(vortex_err!(
39                NotImplemented: "fill_forward",
40                array.encoding()
41            ))
42        })?;
43
44    assert_eq!(
45        filled.len(),
46        array.len(),
47        "FillForward length mismatch {}",
48        array.encoding()
49    );
50    assert_eq!(
51        filled.dtype(),
52        array.dtype(),
53        "FillForward dtype mismatch {}",
54        array.encoding()
55    );
56
57    Ok(filled)
58}