Skip to main content

vortex_array/arrays/varbin/
accessor.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::iter;
5
6use crate::ToCanonical;
7use crate::accessor::ArrayAccessor;
8use crate::arrays::varbin::VarBinArray;
9use crate::match_each_integer_ptype;
10use crate::validity::Validity;
11use crate::vtable::ValidityHelper;
12
13impl ArrayAccessor<[u8]> for VarBinArray {
14    fn with_iterator<F, R>(&self, f: F) -> R
15    where
16        F: for<'a> FnOnce(&mut dyn Iterator<Item = Option<&'a [u8]>>) -> R,
17    {
18        let offsets = self.offsets().to_primitive();
19        let validity = self.validity();
20
21        let bytes = self.bytes();
22        let bytes = bytes.as_slice();
23
24        match_each_integer_ptype!(offsets.ptype(), |T| {
25            let offsets = offsets.as_slice::<T>();
26
27            #[allow(clippy::cast_possible_truncation)]
28            match validity {
29                Validity::NonNullable | Validity::AllValid => {
30                    let mut iter = offsets
31                        .windows(2)
32                        .map(|w| Some(&bytes[w[0] as usize..w[1] as usize]));
33                    f(&mut iter)
34                }
35                Validity::AllInvalid => f(&mut iter::repeat_n(None, self.len())),
36                Validity::Array(v) => {
37                    let validity = v.to_bool().into_bit_buffer();
38                    let mut iter = offsets
39                        .windows(2)
40                        .zip(validity.iter())
41                        .map(|(w, valid)| valid.then(|| &bytes[w[0] as usize..w[1] as usize]));
42                    f(&mut iter)
43                }
44            }
45        })
46    }
47}
48
49impl ArrayAccessor<[u8]> for &VarBinArray {
50    fn with_iterator<F, R>(&self, f: F) -> R
51    where
52        F: for<'a> FnOnce(&mut dyn Iterator<Item = Option<&'a [u8]>>) -> R,
53    {
54        <VarBinArray as ArrayAccessor<[u8]>>::with_iterator(*self, f)
55    }
56}