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}
69
70#[cfg(test)]
71mod tests {
72    use std::sync::Arc;
73
74    use arrow_array::FixedSizeBinaryArray;
75    use arrow_array::RecordBatch;
76    use arrow_array::RecordBatchIterator;
77    use arrow_schema::DataType;
78    use arrow_schema::Field;
79    use arrow_schema::Schema;
80    use arrow_schema::extension::Uuid as ArrowUuid;
81    use vortex_array::array_session;
82    use vortex_array::arrays::Struct;
83    use vortex_array::arrays::struct_::StructArrayExt;
84    use vortex_array::extension::uuid::Uuid;
85    use vortex_error::VortexExpect;
86
87    use super::*;
88    use crate::ArrowSessionExt;
89
90    /// The adapter imports through the [`ArrowSession`], so a UUID column arrives as a Vortex
91    /// extension array rather than its `FixedSizeList` storage.
92    #[test]
93    fn stream_preserves_extension_types() -> VortexResult<()> {
94        let mut field = Field::new("id", DataType::FixedSizeBinary(16), false);
95        field.try_with_extension_type(ArrowUuid)?;
96        let schema = Arc::new(Schema::new(vec![field]));
97        let ids = FixedSizeBinaryArray::try_from_iter(
98            [*b"0123456789abcdef", *b"fedcba9876543210"].into_iter(),
99        )?;
100        let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(ids)])?;
101
102        let reader = RecordBatchIterator::new([Ok(batch)], schema);
103        let stream = ffi_stream::ArrowArrayStreamReader::try_new(
104            ffi_stream::FFI_ArrowArrayStream::new(Box::new(reader)),
105        )?;
106
107        let vortex_session = array_session();
108        let mut adapter = ArrowArrayStreamAdapter::try_new(&vortex_session.arrow(), stream)?;
109
110        let DType::Struct(fields, _) = adapter.dtype().clone() else {
111            panic!("expected a struct dtype, got {}", adapter.dtype());
112        };
113        assert!(
114            fields
115                .field_by_index(0)
116                .vortex_expect("one field")
117                .as_extension()
118                .is::<Uuid>()
119        );
120
121        let array = adapter.next().vortex_expect("one batch")?;
122        assert_eq!(array.dtype(), adapter.dtype());
123        assert!(
124            array
125                .as_::<Struct>()
126                .unmasked_field(0)
127                .dtype()
128                .as_extension()
129                .is::<Uuid>()
130        );
131        assert!(adapter.next().is_none());
132
133        Ok(())
134    }
135}
vortex_arrow/iter.rs

vortex_arrow/
iter.rs
