Arrow iterator conversions (#3875)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/Cargo.toml

@@ -23,7 +23,7 @@ rustversion = { workspace = true }
 arbitrary = { workspace = true, optional = true }
 arcref = { workspace = true }
 arrow-arith = { workspace = true }
-arrow-array = { workspace = true }
+arrow-array = { workspace = true, features = ["ffi"] }
 arrow-buffer = { workspace = true }
 arrow-cast = { workspace = true }
 arrow-data = { workspace = true }

vortex-python/src/record_batch_reader.rs vortex-array/src/arrow/iter.rsvortex-python/src/record_batch_reader.rs renamed to vortex-array/src/arrow/iter.rs

@@ -4,11 +4,47 @@
 use std::sync::Arc;
 
 use arrow_array::cast::AsArray;
-use arrow_array::{RecordBatch, RecordBatchReader};
+use arrow_array::{RecordBatch, RecordBatchReader, ffi_stream};
 use arrow_schema::{ArrowError, DataType, SchemaRef};
-use vortex::arrow::compute::to_arrow;
-use vortex::error::VortexResult;
-use vortex::iter::ArrayIterator;
+use vortex_dtype::DType;
+use vortex_dtype::arrow::FromArrowType;
+use vortex_error::{VortexError, VortexResult};
+
+use crate::ArrayRef;
+use crate::arrow::FromArrowArray;
+use crate::arrow::compute::to_arrow;
+use crate::iter::ArrayIterator;
+
+/// An adapter for converting an `ArrowArrayStreamReader` into a Vortex `ArrayStream`.
+pub struct ArrowArrayStreamAdapter {
+    stream: ffi_stream::ArrowArrayStreamReader,
+    dtype: DType,
+}
+
+impl ArrowArrayStreamAdapter {
+    pub fn new(stream: ffi_stream::ArrowArrayStreamReader, dtype: DType) -> Self {
+        Self { stream, dtype }
+    }
+}
+
+impl ArrayIterator for ArrowArrayStreamAdapter {
+    fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+}
+
+impl Iterator for ArrowArrayStreamAdapter {
+    type Item = VortexResult<ArrayRef>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let batch = self.stream.next()?;
+
+        Some(batch.map_err(VortexError::from).map(|b| {
+            debug_assert_eq!(&self.dtype, &DType::from_arrow(b.schema()));
+            ArrayRef::from_arrow(b, false)
+        }))
+    }
+}
 
 /// Adapter for converting a [`ArrayIterator`] into an Arrow [`RecordBatchReader`].
 pub struct VortexRecordBatchReader<I> {

vortex-array/src/arrow/mod.rs

@@ -11,9 +11,11 @@ mod array;
 pub mod compute;
 mod convert;
 mod datum;
+mod iter;
 mod record_batch;
 
 pub use datum::*;
+pub use iter::*;
 
 use crate::arrow::compute::ToArrowOptions;
 

vortex-python/src/dataset.rs

@@ -9,6 +9,7 @@ use arrow_schema::SchemaRef;
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::types::PyString;
+use vortex::arrow::VortexRecordBatchReader;
 use vortex::dtype::FieldName;
 use vortex::error::VortexResult;
 use vortex::expr::{ExprRef, SelectExpr, root};
@@ -20,7 +21,6 @@ use crate::arrays::PyArrayRef;
 use crate::expr::PyExpr;
 use crate::iter::ArrayStreamToIterator;
 use crate::object_store_urls::object_store_from_url;
-use crate::record_batch_reader::VortexRecordBatchReader;
 use crate::{TOKIO_RUNTIME, install_module};
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {

vortex-python/src/file.rs

@@ -9,6 +9,7 @@ use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::types::PyList;
 use vortex::ToCanonical;
+use vortex::arrow::VortexRecordBatchReader;
 use vortex::compute::cast;
 use vortex::dtype::Nullability::NonNullable;
 use vortex::dtype::{DType, PType};
@@ -24,7 +25,6 @@ use crate::dataset::PyVortexDataset;
 use crate::dtype::PyDType;
 use crate::expr::PyExpr;
 use crate::iter::{ArrayStreamToIterator, PyArrayIterator};
-use crate::record_batch_reader::VortexRecordBatchReader;
 use crate::{TOKIO_RUNTIME, install_module};
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {

vortex-python/src/iter/mod.rs

@@ -13,6 +13,7 @@ use pyo3::prelude::*;
 use pyo3::types::PyIterator;
 use pyo3::{Bound, PyResult, Python};
 pub(crate) use stream::*;
+use vortex::arrow::VortexRecordBatchReader;
 use vortex::dtype::DType;
 use vortex::iter::{ArrayIterator, ArrayIteratorAdapter, ArrayIteratorExt};
 use vortex::{Canonical, IntoArray};
@@ -21,7 +22,6 @@ use crate::arrays::PyArrayRef;
 use crate::dtype::PyDType;
 use crate::install_module;
 use crate::iter::python::PythonArrayIterator;
-use crate::record_batch_reader::VortexRecordBatchReader;
 
 pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
     let m = PyModule::new(py, "iter")?;

vortex-python/src/lib.rs

@@ -19,7 +19,6 @@ mod io;
 mod iter;
 mod object_store_urls;
 mod python_repr;
-mod record_batch_reader;
 mod registry;
 pub(crate) mod scalar;
 mod serde;