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