From 454881a559c3e7ea0b43804e5475c1b7b28159f6 Mon Sep 17 00:00:00 2001 From: Kevin Jacobs Date: Fri, 27 Mar 2026 08:43:51 -0400 Subject: [PATCH 1/6] Accept iterables in RecordBatchReader.from_batches Changes from_batches to accept any Python iterable (list, generator, etc.) instead of requiring a Sequence. Generators are consumed lazily, enabling streaming writes without materializing all batches in memory. Closes #493 Signed-off-by: Kevin Jacobs --- .../arro3/core/_record_batch_reader.pyi | 6 +- pyo3-arrow/src/record_batch_reader.rs | 61 +++++++++++++++---- tests/core/test_ffi.py | 34 +++++++++++ 3 files changed, 88 insertions(+), 13 deletions(-) diff --git a/arro3-core/python/arro3/core/_record_batch_reader.pyi b/arro3-core/python/arro3/core/_record_batch_reader.pyi index 64612105..2c73dec6 100644 --- a/arro3-core/python/arro3/core/_record_batch_reader.pyi +++ b/arro3-core/python/arro3/core/_record_batch_reader.pyi @@ -1,3 +1,4 @@ +from collections.abc import Iterable from typing import Sequence from ._record_batch import RecordBatch @@ -53,13 +54,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 baa223d8..6278f93e 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/pyo3-arrow/src/record_batch_reader.rs @@ -1,8 +1,8 @@ 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::*; @@ -19,6 +19,43 @@ use crate::input::AnyRecordBatch; use crate::schema::display_schema; use crate::{PyRecordBatch, PySchema, PyTable}; +/// A RecordBatchReader that lazily pulls batches from a Python iterator. +/// +/// The `iter` field holds a Python object that has already been converted to +/// an iterator (via `__iter__`). Each call to `next()` acquires the GIL and +/// calls `__next__` on it. +struct PyIteratorRecordBatchReader { + schema: SchemaRef, + iter: Py, +} + +// Safety: Py is Send and only accessed while holding the GIL. +unsafe impl Send for PyIteratorRecordBatchReader {} + +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 +254,17 @@ 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: &Bound, + ) -> PyResult { + let py = batches.py(); + let iter = batches.call_method0(intern!(py, "__iter__"))?; + Ok(Self::new(Box::new(PyIteratorRecordBatchReader { + schema: schema.into_inner(), + iter: iter.unbind(), + }))) } #[classmethod] diff --git a/tests/core/test_ffi.py b/tests/core/test_ffi.py index 25fe7073..50223a79 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]}) From 4bd49b3316399f376d5b2391a62cffa459859c55 Mon Sep 17 00:00:00 2001 From: Kevin Jacobs Date: Fri, 27 Mar 2026 08:56:18 -0400 Subject: [PATCH 2/6] fix: address clippy and ruff lints Signed-off-by: Kevin Jacobs --- arro3-core/python/arro3/core/_record_batch_reader.pyi | 1 - pyo3-arrow/src/record_batch_reader.rs | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/arro3-core/python/arro3/core/_record_batch_reader.pyi b/arro3-core/python/arro3/core/_record_batch_reader.pyi index 2c73dec6..fd6fc506 100644 --- a/arro3-core/python/arro3/core/_record_batch_reader.pyi +++ b/arro3-core/python/arro3/core/_record_batch_reader.pyi @@ -1,5 +1,4 @@ from collections.abc import Iterable -from typing import Sequence from ._record_batch import RecordBatch from ._schema import Schema diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index 6278f93e..cea4ce49 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/pyo3-arrow/src/record_batch_reader.rs @@ -1,7 +1,7 @@ use std::fmt::Display; use std::sync::{Arc, Mutex}; -use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, StructArray}; +use arrow_array::{ArrayRef, RecordBatch, RecordBatchReader, StructArray}; use arrow_schema::{ArrowError, Field, SchemaRef}; use pyo3::exceptions::{PyIOError, PyStopIteration, PyValueError}; use pyo3::intern; From f715c1bdf202bfb78ce52cd71ff70a201aafcdac Mon Sep 17 00:00:00 2001 From: Kevin Jacobs Date: Fri, 27 Mar 2026 10:35:01 -0400 Subject: [PATCH 3/6] fix: fast path for sequences, fix safety comment Use eager extraction for sequences (no per-batch GIL overhead), fall back to lazy iteration for generators and other iterables. Signed-off-by: Kevin Jacobs --- pyo3-arrow/src/record_batch_reader.rs | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index cea4ce49..d0263286 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/pyo3-arrow/src/record_batch_reader.rs @@ -1,7 +1,7 @@ use std::fmt::Display; use std::sync::{Arc, Mutex}; -use arrow_array::{ArrayRef, RecordBatch, RecordBatchReader, StructArray}; +use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, StructArray}; use arrow_schema::{ArrowError, Field, SchemaRef}; use pyo3::exceptions::{PyIOError, PyStopIteration, PyValueError}; use pyo3::intern; @@ -29,7 +29,8 @@ struct PyIteratorRecordBatchReader { iter: Py, } -// Safety: Py is Send and only accessed while holding the GIL. +// Safety: The contained Py is only ever accessed while holding the GIL +// (via Python::attach in Iterator::next), so sending across threads is safe. unsafe impl Send for PyIteratorRecordBatchReader {} impl Iterator for PyIteratorRecordBatchReader { @@ -259,6 +260,20 @@ impl PyRecordBatchReader { schema: PySchema, batches: &Bound, ) -> PyResult { + // Fast path: if batches is a sequence (list, tuple), extract eagerly + // to avoid per-batch GIL acquisition overhead. + if let Ok(vec) = batches.extract::>() { + let batches = vec + .into_iter() + .map(|batch| batch.into_inner()) + .collect::>(); + return Ok(Self::new(Box::new(RecordBatchIterator::new( + batches.into_iter().map(Ok), + schema.into_inner(), + )))); + } + + // Slow path: consume lazily from any iterable (generators, etc.) let py = batches.py(); let iter = batches.call_method0(intern!(py, "__iter__"))?; Ok(Self::new(Box::new(PyIteratorRecordBatchReader { From 26d80ab006eebb6afad835c22fc4b8d22295494f Mon Sep 17 00:00:00 2001 From: Kevin Jacobs Date: Fri, 27 Mar 2026 10:36:57 -0400 Subject: [PATCH 4/6] fix: remove unnecessary unsafe impl Send Py is already Send in PyO3 0.28, so PyIteratorRecordBatchReader auto-derives Send. The Send bound is enforced at the call site by Box. Signed-off-by: Kevin Jacobs --- pyo3-arrow/src/record_batch_reader.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index d0263286..4525f305 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/pyo3-arrow/src/record_batch_reader.rs @@ -29,9 +29,6 @@ struct PyIteratorRecordBatchReader { iter: Py, } -// Safety: The contained Py is only ever accessed while holding the GIL -// (via Python::attach in Iterator::next), so sending across threads is safe. -unsafe impl Send for PyIteratorRecordBatchReader {} impl Iterator for PyIteratorRecordBatchReader { type Item = Result; From 059fefb3baf26108b16980b547dca2dcf0558e40 Mon Sep 17 00:00:00 2001 From: Kevin Jacobs Date: Fri, 27 Mar 2026 10:41:20 -0400 Subject: [PATCH 5/6] refactor: use FromPyObject enum for dispatch Per review feedback, use a RecordBatchInput enum with FromPyObject to dispatch between sequence (eager) and iterator (lazy) paths. Signed-off-by: Kevin Jacobs --- pyo3-arrow/src/record_batch_reader.rs | 65 ++++++++++++++++----------- 1 file changed, 40 insertions(+), 25 deletions(-) diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index 4525f305..2a405eb3 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/pyo3-arrow/src/record_batch_reader.rs @@ -6,6 +6,7 @@ use arrow_schema::{ArrowError, Field, SchemaRef}; use pyo3::exceptions::{PyIOError, PyStopIteration, PyValueError}; use pyo3::intern; use pyo3::prelude::*; +use pyo3::Borrowed; use pyo3::types::{PyCapsule, PyTuple, PyType}; use crate::error::PyArrowResult; @@ -19,17 +20,35 @@ 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. /// -/// The `iter` field holds a Python object that has already been converted to -/// an iterator (via `__iter__`). Each call to `next()` acquires the GIL and -/// calls `__next__` on it. +/// 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; @@ -255,28 +274,24 @@ impl PyRecordBatchReader { fn from_batches( _cls: &Bound, schema: PySchema, - batches: &Bound, - ) -> PyResult { - // Fast path: if batches is a sequence (list, tuple), extract eagerly - // to avoid per-batch GIL acquisition overhead. - if let Ok(vec) = batches.extract::>() { - let batches = vec - .into_iter() - .map(|batch| batch.into_inner()) - .collect::>(); - return Ok(Self::new(Box::new(RecordBatchIterator::new( - batches.into_iter().map(Ok), - schema.into_inner(), - )))); + 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 })) + } } - - // Slow path: consume lazily from any iterable (generators, etc.) - let py = batches.py(); - let iter = batches.call_method0(intern!(py, "__iter__"))?; - Ok(Self::new(Box::new(PyIteratorRecordBatchReader { - schema: schema.into_inner(), - iter: iter.unbind(), - }))) } #[classmethod] From 799b6ac26c8fcbd85d4989805afe324298eb0a1d Mon Sep 17 00:00:00 2001 From: Kevin Jacobs Date: Fri, 27 Mar 2026 10:43:59 -0400 Subject: [PATCH 6/6] style: cargo fmt Signed-off-by: Kevin Jacobs --- pyo3-arrow/src/record_batch_reader.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/pyo3-arrow/src/record_batch_reader.rs b/pyo3-arrow/src/record_batch_reader.rs index 2a405eb3..c25e93ef 100644 --- a/pyo3-arrow/src/record_batch_reader.rs +++ b/pyo3-arrow/src/record_batch_reader.rs @@ -6,8 +6,8 @@ use arrow_schema::{ArrowError, Field, SchemaRef}; use pyo3::exceptions::{PyIOError, PyStopIteration, PyValueError}; use pyo3::intern; use pyo3::prelude::*; -use pyo3::Borrowed; use pyo3::types::{PyCapsule, PyTuple, PyType}; +use pyo3::Borrowed; use crate::error::PyArrowResult; use crate::export::{Arro3RecordBatch, Arro3Schema, Arro3Table}; @@ -271,11 +271,7 @@ impl PyRecordBatchReader { } #[classmethod] - fn from_batches( - _cls: &Bound, - schema: PySchema, - batches: RecordBatchInput, - ) -> Self { + fn from_batches(_cls: &Bound, schema: PySchema, batches: RecordBatchInput) -> Self { let schema = schema.into_inner(); match batches { RecordBatchInput::Sequence(vec) => {