feat: accept iterables in RecordBatchReader.from_batches by kevinjacobs-delfi · Pull Request #495 · kylebarron/arro3

kevinjacobs-delfi · 2026-03-27T12:44:18Z

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.

def batch_gen():
    for chunk in chunks:
        yield RecordBatch(...)

reader = RecordBatchReader.from_batches(schema, batch_gen())
write_parquet(reader, path, ...)

Closes #493

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 kylebarron#493 Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

kylebarron · 2026-03-27T14:10:40Z

+/// 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<PyAny>,
+}


This is a performance regression for any existing case of a sequence of batches, because now it needs to acquire the GIL for every iteration

In a micro-benchmark, absolutely. In real world cases where record batches contain a meaningful amount of data, probably not. If you feel strongly, then I can probably special case Sequence from Iterable to win very microsecond back.

I think acquiring the GIL can have real-world performance overhead and I'd rather not introduce it for the existing approach where it doesn't need it.

You can define an enum like

enum PyRecordBatchReaderInput { Sequence(Vec<PyRecordBatch>) Iterator(PyRecordBatchIterator) }

and then implement FromPyObject on that where first it checks if it's a sequence, and then extracts to the Sequence variant. And falls back to Iterator

No worries. Done.

kylebarron · 2026-03-27T14:11:53Z

+// Safety: Py<PyAny> is Send and only accessed while holding the GIL.
+unsafe impl Send for PyIteratorRecordBatchReader {}


Is it? If Py<PyAny> were Send then I don't think you'd need this unsafe impl.

PyAny is not Send, so Py<PyAny> is not Send

You're right-- sorry. My Rust isn't very advanced. I'll fix.

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 <kevin.jacobs@delfidiagnostics.com>

Py<PyAny> is already Send in PyO3 0.28, so PyIteratorRecordBatchReader auto-derives Send. The Send bound is enforced at the call site by Box<dyn RecordBatchReader + Send>. Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

Per review feedback, use a RecordBatchInput enum with FromPyObject to dispatch between sequence (eager) and iterator (lazy) paths. Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

kevinjacobs-delfi · 2026-03-30T12:33:04Z

Please let me know what else you need before merging this PR. I'm finding arro3 extremely useful and have more small fixes and improvements in the works. Thanks!

kylebarron · 2026-04-15T20:10:24Z

+    Sequence(Vec<PyRecordBatch>),
+    Iterator(Py<PyAny>),


It would be good to double-check here what pyo3 auto-converts to a sequence.

It'll convert anything with an __iter__ and __len__ to a Vec<T>?

Perhaps it makes sense to make a one-off test function just to validate how it handles inputs. We don't want it automatically materializing iterators to a sequence.

kylebarron · 2026-04-15T20:12:08Z

+        if let Ok(vec) = ob.extract::<Vec<PyRecordBatch>>() {
+            Ok(Self::Sequence(vec))
+        } else {
+            let iter = ob.call_method0(intern!(ob.py(), "__iter__"))?;


Here we should have a better error message. If __iter__ doesn't exist, we should print something like "RecordBatchInput requires a sequence or iterator of record batches"

By default the ? will give a somewhat opaque error message of just "__iter__ does not exist" or something like that.

kylebarron · 2026-04-15T20:13:31Z

+            batches: The existing batches. Can be a list, generator, or any
+                iterable. Generators are consumed lazily.


We should clarify; iterators are probably also consumed lazily.

kylebarron · 2026-04-15T20:15:27Z

+            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 }))
+            }
+        }


Nit: I'd move this onto a private method on RecordBatchInput named into_record_batch_iterator which returns a Box<dyn RecordBatchIterator>.

Then this call site can stay clean

kevinjacobs-delfi added 2 commits March 27, 2026 08:43

fix: address clippy and ruff lints

4bd49b3

Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

kevinjacobs-delfi changed the title ~~Accept iterables in RecordBatchReader.from_batches~~ feat: accept iterables in RecordBatchReader.from_batches Mar 27, 2026

github-actions Bot added the feat label Mar 27, 2026

kylebarron reviewed Mar 27, 2026

View reviewed changes

kevinjacobs-delfi added 4 commits March 27, 2026 10:35

fix: fast path for sequences, fix safety comment

f715c1b

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 <kevin.jacobs@delfidiagnostics.com>

fix: remove unnecessary unsafe impl Send

26d80ab

Py<PyAny> is already Send in PyO3 0.28, so PyIteratorRecordBatchReader auto-derives Send. The Send bound is enforced at the call site by Box<dyn RecordBatchReader + Send>. Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

refactor: use FromPyObject enum for dispatch

059fefb

Per review feedback, use a RecordBatchInput enum with FromPyObject to dispatch between sequence (eager) and iterator (lazy) paths. Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

style: cargo fmt

799b6ac

Signed-off-by: Kevin Jacobs <kevin.jacobs@delfidiagnostics.com>

kylebarron reviewed Apr 15, 2026

View reviewed changes

		// Safety: Py<PyAny> is Send and only accessed while holding the GIL.
		unsafe impl Send for PyIteratorRecordBatchReader {}

		batches: The existing batches. Can be a list, generator, or any
		iterable. Generators are consumed lazily.

Conversation

kevinjacobs-delfi commented Mar 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevinjacobs-delfi commented Mar 30, 2026

Uh oh!

kylebarron Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kylebarron Apr 15, 2026 •

edited

Loading