vortex_array/arrays/varbin/
accessor.rs

1use std::iter;
2
3use vortex_dtype::match_each_integer_ptype;
4use vortex_error::VortexResult;
5
6use crate::accessor::ArrayAccessor;
7use crate::arrays::varbin::VarBinArray;
8use crate::validity::Validity;
9use crate::variants::PrimitiveArrayTrait;
10use crate::{Array, ToCanonical};
11
12impl ArrayAccessor<[u8]> for VarBinArray {
13    fn with_iterator<F, R>(&self, f: F) -> VortexResult<R>
14    where
15        F: for<'a> FnOnce(&mut (dyn Iterator<Item = Option<&'a [u8]>>)) -> R,
16    {
17        let offsets = self.offsets().to_primitive()?;
18        let validity = self.validity();
19
20        let bytes = self.bytes();
21        let bytes = bytes.as_slice();
22
23        match_each_integer_ptype!(offsets.ptype(), |$T| {
24            let offsets = offsets.as_slice::<$T>();
25
26            match validity {
27                Validity::NonNullable | Validity::AllValid => {
28                    let mut iter = offsets
29                        .windows(2)
30                        .map(|w| Some(&bytes[w[0] as usize..w[1] as usize]));
31                    Ok(f(&mut iter))
32                }
33                Validity::AllInvalid => Ok(f(&mut iter::repeat_n(None, self.len()))),
34                Validity::Array(v) => {
35                    let validity = v.to_bool()?;
36                    let mut iter = offsets
37                        .windows(2)
38                        .zip(validity.boolean_buffer())
39                        .map(|(w, valid)| {
40                            if valid {
41                                Some(&bytes[w[0] as usize..w[1] as usize])
42                            } else {
43                                None
44                            }
45                        });
46                    Ok(f(&mut iter))
47                }
48            }
49        })
50    }
51}