vortex_arrow/
iter.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_array::RecordBatchReader;
5use arrow_array::ffi_stream;
6use arrow_schema::SchemaRef;
7use vortex_array::ArrayRef;
8use vortex_array::dtype::DType;
9use vortex_array::iter::ArrayIterator;
10use vortex_error::VortexError;
11use vortex_error::VortexResult;
12
13use crate::ArrowSession;
14
15/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayIterator`.
16///
17/// The stream's schema and batches are imported through the [`ArrowSession`], so Arrow extension
18/// types are routed to their registered import plugins instead of being flattened into their
19/// storage types.
20pub struct ArrowArrayStreamAdapter {
21    stream: ffi_stream::ArrowArrayStreamReader,
22    session: ArrowSession,
23    schema: SchemaRef,
24    dtype: DType,
25}
26
27impl ArrowArrayStreamAdapter {
28    /// Adapt `stream`, importing its schema and each of its batches through `session`.
29    ///
30    /// The adapter holds a clone of `session`, which shares the plugin registries with it, so
31    /// plugins registered after construction are still observed.
32    ///
33    /// The schema declared by the stream is the authoritative schema for the import: every batch
34    /// is converted against it, so extension metadata survives even if an individual batch has
35    /// lost it.
36    pub fn try_new(
37        session: &ArrowSession,
38        stream: ffi_stream::ArrowArrayStreamReader,
39    ) -> VortexResult<Self> {
40        let schema = stream.schema();
41        let dtype = session.from_arrow_schema(schema.as_ref())?;
42        Ok(Self {
43            stream,
44            session: session.clone(),
45            schema,
46            dtype,
47        })
48    }
49}
50
51impl ArrayIterator for ArrowArrayStreamAdapter {
52    fn dtype(&self) -> &DType {
53        &self.dtype
54    }
55}
56
57impl Iterator for ArrowArrayStreamAdapter {
58    type Item = VortexResult<ArrayRef>;
59
60    fn next(&mut self) -> Option<Self::Item> {
61        let batch = self.stream.next()?;
62
63        Some(batch.map_err(VortexError::from).and_then(|b| {
64            self.session
65                .from_arrow_record_batch(b, self.schema.as_ref())
66        }))
67    }
68}
vortex_arrow/iter.rs

vortex_arrow/
iter.rs
