Skip to main content

polars_arrow/io/
iterator.rs

1pub use streaming_iterator::StreamingIterator;
2
3/// A [`StreamingIterator`] with an internal buffer of [`Vec<u8>`] used to efficiently
4/// present items of type `T` as `&[u8]`.
5/// It is generic over the type `T` and the transformation `F: T -> &[u8]`.
6pub struct BufStreamingIterator<I, F, T>
7where
8    I: Iterator<Item = T>,
9    F: FnMut(T, &mut Vec<u8>),
10{
11    iterator: I,
12    f: F,
13    buffer: Vec<u8>,
14    is_valid: bool,
15}
16
17impl<I, F, T> BufStreamingIterator<I, F, T>
18where
19    I: Iterator<Item = T>,
20    F: FnMut(T, &mut Vec<u8>),
21{
22    #[inline]
23    pub fn new(iterator: I, f: F, buffer: Vec<u8>) -> Self {
24        Self {
25            iterator,
26            f,
27            buffer,
28            is_valid: false,
29        }
30    }
31
32    pub fn into_inner(self) -> (I, F) {
33        (self.iterator, self.f)
34    }
35}
36
37impl<I, F, T> StreamingIterator for BufStreamingIterator<I, F, T>
38where
39    I: Iterator<Item = T>,
40    F: FnMut(T, &mut Vec<u8>),
41{
42    type Item = [u8];
43
44    #[inline]
45    fn advance(&mut self) {
46        let a = self.iterator.next();
47        if let Some(a) = a {
48            self.is_valid = true;
49            self.buffer.clear();
50            (self.f)(a, &mut self.buffer);
51        } else {
52            self.is_valid = false;
53        }
54    }
55
56    #[inline]
57    fn get(&self) -> Option<&Self::Item> {
58        if self.is_valid {
59            Some(&self.buffer)
60        } else {
61            None
62        }
63    }
64
65    #[inline]
66    fn size_hint(&self) -> (usize, Option<usize>) {
67        self.iterator.size_hint()
68    }
69}