vortex_array/arrays/bool/compute/
fill_forward.rs

1use arrow_buffer::BooleanBuffer;
2use itertools::Itertools;
3use vortex_dtype::Nullability;
4use vortex_error::VortexResult;
5use vortex_mask::AllOr;
6
7use crate::arrays::{BoolArray, BoolEncoding};
8use crate::compute::FillForwardFn;
9use crate::validity::Validity;
10use crate::{Array, ArrayRef};
11
12impl FillForwardFn<&BoolArray> for BoolEncoding {
13    fn fill_forward(&self, array: &BoolArray) -> VortexResult<ArrayRef> {
14        let validity = array.validity_mask()?;
15
16        // nothing to see or do in this case
17        if array.dtype().nullability() == Nullability::NonNullable {
18            return Ok(array.to_array().into_array());
19        }
20
21        match validity.boolean_buffer() {
22            AllOr::All => {
23                // all valid, but we need to convert to non-nullable
24                Ok(BoolArray::new(array.boolean_buffer().clone(), Validity::AllValid).into_array())
25            }
26            AllOr::None => {
27                // all invalid => fill with default value (false)
28                Ok(
29                    BoolArray::new(BooleanBuffer::new_unset(array.len()), Validity::AllValid)
30                        .into_array(),
31                )
32            }
33            AllOr::Some(validity) => {
34                let bools = array.boolean_buffer();
35                let mut last_value = false;
36                let buffer = BooleanBuffer::from_iter(bools.iter().zip_eq(validity.iter()).map(
37                    |(v, valid)| {
38                        if valid {
39                            last_value = v;
40                        }
41                        last_value
42                    },
43                ));
44                Ok(BoolArray::new(buffer, Validity::AllValid).into_array())
45            }
46        }
47    }
48}
49
50#[cfg(test)]
51mod test {
52    use crate::arrays::BoolArray;
53    use crate::validity::Validity;
54    use crate::{ToCanonical, compute};
55
56    #[test]
57    fn fill_forward() {
58        let barr = BoolArray::from_iter(vec![None, Some(false), None, Some(true), None]);
59        let filled_bool = compute::fill_forward(&barr).unwrap().to_bool().unwrap();
60        assert_eq!(
61            filled_bool.boolean_buffer().iter().collect::<Vec<bool>>(),
62            vec![false, false, false, true, true]
63        );
64        assert_eq!(filled_bool.validity(), &Validity::AllValid);
65    }
66}
vortex_array/arrays/bool/compute/fill_forward.rs

vortex_array/arrays/bool/compute/
fill_forward.rs
