vortex_array/arrow/
iter.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_array::ffi_stream;
5use vortex_dtype::DType;
6use vortex_dtype::arrow::FromArrowType;
7use vortex_error::{VortexError, VortexResult};
8
9use crate::ArrayRef;
10use crate::arrow::FromArrowArray;
11use crate::iter::ArrayIterator;
12
13/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`.
14pub struct ArrowArrayStreamAdapter {
15    stream: ffi_stream::ArrowArrayStreamReader,
16    dtype: DType,
17}
18
19impl ArrowArrayStreamAdapter {
20    pub fn new(stream: ffi_stream::ArrowArrayStreamReader, dtype: DType) -> Self {
21        Self { stream, dtype }
22    }
23}
24
25impl ArrayIterator for ArrowArrayStreamAdapter {
26    fn dtype(&self) -> &DType {
27        &self.dtype
28    }
29}
30
31impl Iterator for ArrowArrayStreamAdapter {
32    type Item = VortexResult<ArrayRef>;
33
34    fn next(&mut self) -> Option<Self::Item> {
35        let batch = self.stream.next()?;
36
37        Some(batch.map_err(VortexError::from).map(|b| {
38            debug_assert_eq!(&self.dtype, &DType::from_arrow(b.schema()));
39            ArrayRef::from_arrow(b, false)
40        }))
41    }
42}
vortex_array/arrow/iter.rs

vortex_array/arrow/
iter.rs
