feat: accept iterables in RecordBatchReader.from_batches

arro3-core/python/arro3/core/_record_batch_reader.pyi

@@ -1,4 +1,4 @@
-from typing import Sequence
+from collections.abc import Iterable
 
 from ._record_batch import RecordBatch
 from ._schema import Schema
@@ -53,13 +53,14 @@ class RecordBatchReader:
         """Construct this object from a bare Arrow PyCapsule"""
     @classmethod
     def from_batches(
-        cls, schema: ArrowSchemaExportable, batches: Sequence[ArrowArrayExportable]
+        cls, schema: ArrowSchemaExportable, batches: Iterable[ArrowArrayExportable]
     ) -> RecordBatchReader:
         """Construct a new RecordBatchReader from existing data.
 
         Args:
             schema: The schema of the Arrow batches.
-            batches: The existing batches.
+            batches: The existing batches. Can be a list, generator, or any
+                iterable. Generators are consumed lazily.
         """
     @classmethod
     def from_stream(cls, data: ArrowStreamExportable) -> RecordBatchReader:

pyo3-arrow/src/record_batch_reader.rs

@@ -1,12 +1,13 @@
 use std::fmt::Display;
 use std::sync::{Arc, Mutex};
 
-use arrow_array::{ArrayRef, RecordBatchIterator, RecordBatchReader, StructArray};
-use arrow_schema::{Field, SchemaRef};
+use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, StructArray};
+use arrow_schema::{ArrowError, Field, SchemaRef};
 use pyo3::exceptions::{PyIOError, PyStopIteration, PyValueError};
 use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyTuple, PyType};
+use pyo3::Borrowed;
 
 use crate::error::PyArrowResult;
 use crate::export::{Arro3RecordBatch, Arro3Schema, Arro3Table};
@@ -19,6 +20,59 @@ use crate::input::AnyRecordBatch;
 use crate::schema::display_schema;
 use crate::{PyRecordBatch, PySchema, PyTable};
 
+/// Input for `from_batches`: either a sequence (extracted eagerly) or an
+/// arbitrary iterable (consumed lazily via GIL-acquiring iterator).
+enum RecordBatchInput {
+    Sequence(Vec<PyRecordBatch>),
+    Iterator(Py<PyAny>),
+}
+
+impl<'a, 'py> FromPyObject<'a, 'py> for RecordBatchInput {
+    type Error = PyErr;
+
+    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
+        if let Ok(vec) = ob.extract::<Vec<PyRecordBatch>>() {
+            Ok(Self::Sequence(vec))
+        } else {
+            let iter = ob.call_method0(intern!(ob.py(), "__iter__"))?;
+            Ok(Self::Iterator(iter.unbind()))
+        }
+    }
+}
+
+/// A RecordBatchReader that lazily pulls batches from a Python iterator.
+///
+/// Each call to `next()` acquires the GIL and calls `__next__` on the
+/// underlying Python iterator.
+struct PyIteratorRecordBatchReader {
+    schema: SchemaRef,
+    iter: Py<PyAny>,
+}
+
+impl Iterator for PyIteratorRecordBatchReader {
+    type Item = Result<RecordBatch, ArrowError>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        Python::attach(|py| {
+            let iter = self.iter.bind(py);
+            match iter.call_method0(intern!(py, "__next__")) {
+                Ok(obj) => match obj.extract::<PyRecordBatch>() {
+                    Ok(batch) => Some(Ok(batch.into_inner())),
+                    Err(e) => Some(Err(ArrowError::ExternalError(Box::new(e)))),
+                },
+                Err(e) if e.is_instance_of::<PyStopIteration>(py) => None,
+                Err(e) => Some(Err(ArrowError::ExternalError(Box::new(e)))),
+            }
+        })
+    }
+}
+
+impl RecordBatchReader for PyIteratorRecordBatchReader {
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+}
+
 /// A Python-facing Arrow record batch reader.
 ///
 /// This is a wrapper around a [RecordBatchReader].
@@ -217,15 +271,23 @@ impl PyRecordBatchReader {
     }
 
     #[classmethod]
-    fn from_batches(_cls: &Bound<PyType>, schema: PySchema, batches: Vec<PyRecordBatch>) -> Self {
-        let batches = batches
-            .into_iter()
-            .map(|batch| batch.into_inner())
-            .collect::<Vec<_>>();
-        Self::new(Box::new(RecordBatchIterator::new(
-            batches.into_iter().map(Ok),
-            schema.into_inner(),
-        )))
+    fn from_batches(_cls: &Bound<PyType>, schema: PySchema, batches: RecordBatchInput) -> Self {
+        let schema = schema.into_inner();
+        match batches {
+            RecordBatchInput::Sequence(vec) => {
+                let batches = vec
+                    .into_iter()
+                    .map(|batch| batch.into_inner())
+                    .collect::<Vec<_>>();
+                Self::new(Box::new(RecordBatchIterator::new(
+                    batches.into_iter().map(Ok),
+                    schema,
+                )))
+            }
+            RecordBatchInput::Iterator(iter) => {
+                Self::new(Box::new(PyIteratorRecordBatchReader { schema, iter }))
+            }
+        }
     }
 
     #[classmethod]

tests/core/test_ffi.py

@@ -61,6 +61,40 @@ def test_table_metadata_preserved():
     assert pa_table_retour.schema.metadata == metadata
 
 
+def test_record_batch_reader_from_batches_generator():
+    """from_batches accepts a generator and consumes it lazily."""
+    a = pa.array([1, 2, 3], type=pa.int32())
+    table = Table.from_pydict({"a": a})
+    schema = table.schema
+    batches = table.to_batches()
+
+    consumed = []
+
+    def batch_gen():
+        for batch in batches:
+            consumed.append(True)
+            yield batch
+
+    gen = batch_gen()
+    reader = RecordBatchReader.from_batches(schema, gen)
+
+    # Generator not consumed yet
+    assert len(consumed) == 0
+
+    result = reader.read_all()
+    assert result.num_rows == 3
+    assert len(consumed) == 1  # consumed lazily
+
+
+def test_record_batch_reader_from_batches_list():
+    """from_batches still accepts a list (backwards compat)."""
+    a = pa.array([1, 2, 3], type=pa.int32())
+    table = Table.from_pydict({"a": a})
+    reader = RecordBatchReader.from_batches(table.schema, table.to_batches())
+    result = reader.read_all()
+    assert result.num_rows == 3
+
+
 def test_record_batch_reader_metadata_preserved():
     metadata = {b"hello": b"world"}
     pa_table = pa.table({"a": [1, 2, 3]})