pyo3_arrow/ffi/to_python/
chunked.rs

1use arrow_array::ArrayRef;
2use arrow_schema::{ArrowError, FieldRef};
3
4/// Trait for types that can read `ArrayRef`'s.
5///
6/// To create from an iterator, see [ArrayIterator].
7pub trait ArrayReader: Iterator<Item = Result<ArrayRef, ArrowError>> {
8    /// Returns the field of this `ArrayReader`.
9    ///
10    /// Implementation of this trait should guarantee that all `ArrayRef`'s returned by this
11    /// reader should have the same field as returned from this method.
12    fn field(&self) -> FieldRef;
13}
14
15impl<R: ArrayReader + ?Sized> ArrayReader for Box<R> {
16    fn field(&self) -> FieldRef {
17        self.as_ref().field()
18    }
19}
20
21/// An iterator of [`ArrayRef`] with an attached [`FieldRef`]
22pub struct ArrayIterator<I>
23where
24    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
25{
26    inner: I::IntoIter,
27    inner_field: FieldRef,
28}
29
30impl<I> ArrayIterator<I>
31where
32    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
33{
34    /// Create a new [ArrayIterator].
35    ///
36    /// If `iter` is an infallible iterator, use `.map(Ok)`.
37    pub fn new(iter: I, field: FieldRef) -> Self {
38        Self {
39            inner: iter.into_iter(),
40            inner_field: field,
41        }
42    }
43}
44
45impl<I> Iterator for ArrayIterator<I>
46where
47    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
48{
49    type Item = I::Item;
50
51    fn next(&mut self) -> Option<Self::Item> {
52        self.inner.next()
53    }
54
55    fn size_hint(&self) -> (usize, Option<usize>) {
56        self.inner.size_hint()
57    }
58}
59
60impl<I> ArrayReader for ArrayIterator<I>
61where
62    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
63{
64    fn field(&self) -> FieldRef {
65        self.inner_field.clone()
66    }
67}