arrow_array/
array_reader.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Traits for reading a stream of [`ArrayRef`] that share a single [`FieldRef`].
19
20use arrow_schema::{ArrowError, FieldRef};
21
22use crate::array::ArrayRef;
23
24/// Trait for types that can read `ArrayRef`'s.
25///
26/// This is the array-level counterpart to [`RecordBatchReader`], for streams whose
27/// elements are not record batches.
28///
29/// To create from an iterator, see [ArrayIterator]. To read one from the C Stream Interface,
30/// see [`ffi_stream::ArrayStreamReader`].
31///
32/// [`RecordBatchReader`]: crate::RecordBatchReader
33/// [`ffi_stream::ArrayStreamReader`]: crate::ffi_stream::ArrayStreamReader
34pub trait ArrayReader: Iterator<Item = Result<ArrayRef, ArrowError>> {
35    /// Returns the field of this `ArrayReader`.
36    ///
37    /// Implementation of this trait should guarantee that all `ArrayRef`'s returned by this
38    /// reader have the same field as returned from this method.
39    fn field(&self) -> &FieldRef;
40}
41
42impl<R: ArrayReader + ?Sized> ArrayReader for Box<R> {
43    fn field(&self) -> &FieldRef {
44        self.as_ref().field()
45    }
46}
47
48/// Generic implementation of [ArrayReader] that wraps an iterator.
49///
50/// # Example
51///
52/// ```
53/// # use std::sync::Arc;
54/// # use arrow_array::{ArrayIterator, ArrayReader, ArrayRef, Int32Array};
55/// # use arrow_schema::{DataType, Field};
56/// #
57/// let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
58/// let b: ArrayRef = Arc::new(Int32Array::from(vec![3, 4]));
59///
60/// let field = Arc::new(Field::new("values", DataType::Int32, false));
61/// let iter = vec![a.clone(), b.clone()].into_iter().map(Ok);
62/// let mut reader = ArrayIterator::new(iter, field.clone());
63///
64/// assert_eq!(reader.field(), &field);
65/// assert_eq!(&reader.next().unwrap().unwrap(), &a);
66/// assert_eq!(&reader.next().unwrap().unwrap(), &b);
67/// assert!(reader.next().is_none());
68/// ```
69pub struct ArrayIterator<I>
70where
71    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
72{
73    inner: I::IntoIter,
74    inner_field: FieldRef,
75}
76
77impl<I> ArrayIterator<I>
78where
79    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
80{
81    /// Create a new [ArrayIterator].
82    ///
83    /// If `iter` is an infallible iterator, use `.map(Ok)`.
84    pub fn new(iter: I, field: FieldRef) -> Self {
85        Self {
86            inner: iter.into_iter(),
87            inner_field: field,
88        }
89    }
90}
91
92impl<I> Iterator for ArrayIterator<I>
93where
94    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
95{
96    type Item = I::Item;
97
98    fn next(&mut self) -> Option<Self::Item> {
99        self.inner.next()
100    }
101
102    fn size_hint(&self) -> (usize, Option<usize>) {
103        self.inner.size_hint()
104    }
105}
106
107impl<I> ArrayReader for ArrayIterator<I>
108where
109    I: IntoIterator<Item = Result<ArrayRef, ArrowError>>,
110{
111    fn field(&self) -> &FieldRef {
112        &self.inner_field
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use std::sync::Arc;
119
120    use arrow_schema::{ArrowError, DataType, Field};
121
122    use crate::array::{ArrayRef, Int32Array};
123    use crate::{ArrayIterator, ArrayReader};
124
125    fn test_field() -> Arc<Field> {
126        Arc::new(Field::new("values", DataType::Int32, true))
127    }
128
129    #[test]
130    fn array_iterator_reports_its_field_and_yields_its_arrays() {
131        let field = test_field();
132        let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2]));
133        let b: ArrayRef = Arc::new(Int32Array::from(vec![3]));
134
135        let mut reader = ArrayIterator::new(
136            vec![a.clone(), b.clone()].into_iter().map(Ok),
137            field.clone(),
138        );
139
140        assert_eq!(reader.field(), &field);
141        assert_eq!(&reader.next().unwrap().unwrap(), &a);
142        assert_eq!(&reader.next().unwrap().unwrap(), &b);
143        assert!(reader.next().is_none());
144    }
145
146    #[test]
147    fn array_iterator_delegates_size_hint() {
148        let arrays: Vec<Result<ArrayRef, ArrowError>> =
149            vec![Ok(Arc::new(Int32Array::from(vec![1])) as ArrayRef)];
150        let reader = ArrayIterator::new(arrays, test_field());
151
152        assert_eq!(reader.size_hint(), (1, Some(1)));
153    }
154
155    #[test]
156    fn boxed_array_reader_forwards_field() {
157        let field = test_field();
158        let reader = ArrayIterator::new(std::iter::empty(), field.clone());
159
160        let boxed: Box<dyn ArrayReader + Send> = Box::new(reader);
161        assert_eq!(boxed.field(), &field);
162    }
163}
arrow_array/array_reader.rs

arrow_array/
array_reader.rs
