arrow_array/
array_reader.rs

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

arrow_array/
array_reader.rs
