vortex_array/array/bool/compute/
fill_forward.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
use arrow_buffer::BooleanBuffer;
use vortex_dtype::Nullability;
use vortex_error::{vortex_err, VortexResult};

use crate::array::{BoolArray, BoolEncoding};
use crate::compute::FillForwardFn;
use crate::validity::{ArrayValidity, Validity};
use crate::{ArrayDType, ArrayData, ArrayLen, IntoArrayData, ToArrayData};

impl FillForwardFn<BoolArray> for BoolEncoding {
    fn fill_forward(&self, array: &BoolArray) -> VortexResult<ArrayData> {
        let validity = array.logical_validity();
        // nothing to see or do in this case
        if array.dtype().nullability() == Nullability::NonNullable {
            return Ok(array.to_array());
        }

        // all valid, but we need to convert to non-nullable
        if validity.all_valid() {
            return Ok(BoolArray::new(array.boolean_buffer(), Nullability::Nullable).into_array());
        }
        // all invalid => fill with default value (false)
        if validity.all_invalid() {
            return Ok(BoolArray::try_new(
                BooleanBuffer::new_unset(array.len()),
                Validity::AllValid,
            )?
            .into_array());
        }

        let validity = validity
            .to_null_buffer()?
            .ok_or_else(|| vortex_err!("Failed to convert array validity to null buffer"))?;

        let bools = array.boolean_buffer();
        let mut last_value = false;
        let buffer = BooleanBuffer::from_iter(bools.iter().zip(validity.inner().iter()).map(
            |(v, valid)| {
                if valid {
                    last_value = v;
                }
                last_value
            },
        ));
        Ok(BoolArray::try_new(buffer, Validity::AllValid)?.into_array())
    }
}

#[cfg(test)]
mod test {
    use crate::array::BoolArray;
    use crate::validity::Validity;
    use crate::{compute, IntoArrayData};

    #[test]
    fn fill_forward() {
        let barr =
            BoolArray::from_iter(vec![None, Some(false), None, Some(true), None]).into_array();
        let filled_bool = BoolArray::try_from(compute::fill_forward(&barr).unwrap()).unwrap();
        assert_eq!(
            filled_bool.boolean_buffer().iter().collect::<Vec<bool>>(),
            vec![false, false, false, true, true]
        );
        assert_eq!(filled_bool.validity(), Validity::AllValid);
    }
}