Skip to main content

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