vortex_array/arrays/primitive/compute/
fill.rs

1use vortex_buffer::Buffer;
2use vortex_dtype::{Nullability, match_each_native_ptype};
3use vortex_error::VortexResult;
4use vortex_mask::AllOr;
5use vortex_scalar::Scalar;
6
7use crate::arrays::primitive::PrimitiveArray;
8use crate::arrays::{ConstantArray, PrimitiveEncoding};
9use crate::compute::FillForwardFn;
10use crate::validity::Validity;
11use crate::variants::PrimitiveArrayTrait;
12use crate::{Array, ArrayRef};
13
14impl FillForwardFn<&PrimitiveArray> for PrimitiveEncoding {
15    fn fill_forward(&self, array: &PrimitiveArray) -> VortexResult<ArrayRef> {
16        if array.dtype().nullability() == Nullability::NonNullable {
17            return Ok(array.to_array().into_array());
18        }
19
20        match array.validity_mask()?.boolean_buffer() {
21            AllOr::All => Ok(PrimitiveArray::from_byte_buffer(
22                array.byte_buffer().clone(),
23                array.ptype(),
24                Validity::AllValid,
25            )
26            .into_array()),
27            AllOr::None => {
28                match_each_native_ptype!(array.ptype(), |$T| {
29                    let fill_value = Scalar::from($T::default()).cast(array.dtype())?;
30                    return Ok(ConstantArray::new(fill_value, array.len()).into_array())
31                })
32            }
33            AllOr::Some(validity) => {
34                // TODO(ngates): when we take PrimitiveArray by value, we should mutate in-place
35                match_each_native_ptype!(array.ptype(), |$T| {
36                    let as_slice = array.as_slice::<$T>();
37                    let mut last_value = $T::default();
38                    let filled = Buffer::from_iter(
39                        as_slice
40                            .iter()
41                            .zip(validity.into_iter())
42                            .map(|(v, valid)| {
43                                if valid {
44                                    last_value = *v;
45                                }
46                                last_value
47                            })
48                    );
49                    Ok(PrimitiveArray::new(filled, Validity::AllValid).into_array())
50                })
51            }
52        }
53    }
54}
55
56#[cfg(test)]
57mod test {
58    use vortex_buffer::buffer;
59
60    use crate::array::Array;
61    use crate::arrays::BoolArray;
62    use crate::arrays::primitive::PrimitiveArray;
63    use crate::canonical::ToCanonical;
64    use crate::compute::fill_forward;
65    use crate::validity::Validity;
66
67    #[test]
68    fn leading_none() {
69        let arr = PrimitiveArray::from_option_iter([None, Some(8u8), None, Some(10), None]);
70        let p = fill_forward(&arr).unwrap().to_primitive().unwrap();
71        assert_eq!(p.as_slice::<u8>(), vec![0, 8, 8, 10, 10]);
72        assert!(p.validity_mask().unwrap().all_true());
73    }
74
75    #[test]
76    fn all_none() {
77        let arr = PrimitiveArray::from_option_iter([Option::<u8>::None, None, None, None, None]);
78
79        let p = fill_forward(&arr).unwrap().to_primitive().unwrap();
80        assert_eq!(p.as_slice::<u8>(), vec![0, 0, 0, 0, 0]);
81        assert!(p.validity_mask().unwrap().all_true());
82    }
83
84    #[test]
85    fn nullable_non_null() {
86        let arr = PrimitiveArray::new(
87            buffer![8u8, 10, 12, 14, 16],
88            Validity::Array(BoolArray::from_iter([true, true, true, true, true]).into_array()),
89        );
90        let p = fill_forward(&arr).unwrap().to_primitive().unwrap();
91        assert_eq!(p.as_slice::<u8>(), vec![8u8, 10, 12, 14, 16]);
92        assert!(p.validity_mask().unwrap().all_true());
93    }
94}
vortex_array/arrays/primitive/compute/fill.rs

vortex_array/arrays/primitive/compute/
fill.rs
