diff --git a/arro3-core/python/arro3/core/_record_batch_reader.pyi b/arro3-core/python/arro3/core/_record_batch_reader.pyi index 6461210..fd6fc50 100644 --- a/arro3-core/python/arro3/core/_record_batch_reader.pyi +++ b/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: diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index baa223d..c25e93e 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/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), + Iterator(Py), +} + +impl<'a, 'py> FromPyObject<'a, 'py> for RecordBatchInput { + type Error = PyErr; + + fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult { + if let Ok(vec) = ob.extract::>() { + 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, +} + +impl Iterator for PyIteratorRecordBatchReader { + type Item = Result; + + fn next(&mut self) -> Option { + Python::attach(|py| { + let iter = self.iter.bind(py); + match iter.call_method0(intern!(py, "__next__")) { + Ok(obj) => match obj.extract::() { + Ok(batch) => Some(Ok(batch.into_inner())), + Err(e) => Some(Err(ArrowError::ExternalError(Box::new(e)))), + }, + Err(e) if e.is_instance_of::(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, schema: PySchema, batches: Vec) -> Self { - let batches = batches - .into_iter() - .map(|batch| batch.into_inner()) - .collect::>(); - Self::new(Box::new(RecordBatchIterator::new( - batches.into_iter().map(Ok), - schema.into_inner(), - ))) + fn from_batches(_cls: &Bound, 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::>(); + Self::new(Box::new(RecordBatchIterator::new( + batches.into_iter().map(Ok), + schema, + ))) + } + RecordBatchInput::Iterator(iter) => { + Self::new(Box::new(PyIteratorRecordBatchReader { schema, iter })) + } + } } #[classmethod] diff --git a/tests/core/test_ffi.py b/tests/core/test_ffi.py index 25fe707..50223a7 100644 --- a/tests/core/test_ffi.py +++ b/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]})