This is an automated email from the ASF dual-hosted git repository.
zanmato1984 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 7128c9c262 GH-50801: [Python] Expose the `record_batch_reader_source`
Acero node (RecordBatchReaderSourceNodeOptions) (#50802)
7128c9c262 is described below
commit 7128c9c262b5da051e898839f1128185ce050386
Author: You-Cheng Lin <[email protected]>
AuthorDate: Thu Aug 13 10:30:06 2026 +0800
GH-50801: [Python] Expose the `record_batch_reader_source` Acero node
(RecordBatchReaderSourceNodeOptions) (#50802)
### Rationale for this change
Closes #50801
As title, this exposes RecordBatchReaderSourceNodeOptions to Python. With
it, a hash join can build its hash table once and stream the probe side through
as a generator-backed RecordBatchReader, releasing each input chunk as it is
consumed, instead of materializing the full probe table up front (table_source)
or pinning all fragments for the plan's lifetime (dataset scan).
### Are these changes tested?
Yes
### Are there any user-facing changes?
Yeah, users can pass RecordBatchReaderSourceNodeOptions into Declaration,
but its not a breaking change
* GitHub Issue: #50801
Lead-authored-by: You-Cheng Lin
<[email protected]>
Co-authored-by: You-Cheng Lin <[email protected]>
Signed-off-by: Rossi Sun <[email protected]>
---
docs/source/python/api/acero.rst | 1 +
python/pyarrow/_acero.pyx | 25 ++++++++
python/pyarrow/acero.py | 1 +
python/pyarrow/includes/libarrow_acero.pxd | 3 +
python/pyarrow/tests/test_acero.py | 91 ++++++++++++++++++++++++++++++
5 files changed, 121 insertions(+)
diff --git a/docs/source/python/api/acero.rst b/docs/source/python/api/acero.rst
index b75fc33061..363353bf9f 100644
--- a/docs/source/python/api/acero.rst
+++ b/docs/source/python/api/acero.rst
@@ -36,6 +36,7 @@ and to execute this efficiently in a batched manner.
Declaration
ExecNodeOptions
TableSourceNodeOptions
+ RecordBatchReaderSourceNodeOptions
ScanNodeOptions
FilterNodeOptions
ProjectNodeOptions
diff --git a/python/pyarrow/_acero.pyx b/python/pyarrow/_acero.pyx
index 6e88fc7210..e27509469e 100644
--- a/python/pyarrow/_acero.pyx
+++ b/python/pyarrow/_acero.pyx
@@ -79,6 +79,31 @@ class TableSourceNodeOptions(_TableSourceNodeOptions):
self._set_options(table)
+cdef class _RecordBatchReaderSourceNodeOptions(ExecNodeOptions):
+
+ def _set_options(self, RecordBatchReader reader):
+ self.wrapped.reset(
+ new CRecordBatchReaderSourceNodeOptions(reader.reader)
+ )
+
+
+class RecordBatchReaderSourceNodeOptions(_RecordBatchReaderSourceNodeOptions):
+ """
+ A Source node which streams data from a RecordBatchReader.
+
+ This is the option class for the "record_batch_reader_source" node
+ factory.
+
+ Parameters
+ ----------
+ reader : pyarrow.RecordBatchReader
+ The reader which acts as the data source.
+ """
+
+ def __init__(self, RecordBatchReader reader not None):
+ self._set_options(reader)
+
+
cdef class _FilterNodeOptions(ExecNodeOptions):
def _set_options(self, Expression filter_expression not None):
diff --git a/python/pyarrow/acero.py b/python/pyarrow/acero.py
index e475e8db5c..55f03a8bb7 100644
--- a/python/pyarrow/acero.py
+++ b/python/pyarrow/acero.py
@@ -30,6 +30,7 @@ try:
Declaration,
ExecNodeOptions,
TableSourceNodeOptions,
+ RecordBatchReaderSourceNodeOptions,
FilterNodeOptions,
ProjectNodeOptions,
AggregateNodeOptions,
diff --git a/python/pyarrow/includes/libarrow_acero.pxd
b/python/pyarrow/includes/libarrow_acero.pxd
index dc9babee19..7e81a39368 100644
--- a/python/pyarrow/includes/libarrow_acero.pxd
+++ b/python/pyarrow/includes/libarrow_acero.pxd
@@ -42,6 +42,9 @@ cdef extern from "arrow/acero/options.h" namespace
"arrow::acero" nogil:
CTableSourceNodeOptions(shared_ptr[CTable] table)
CTableSourceNodeOptions(shared_ptr[CTable] table, int64_t
max_batch_size)
+ cdef cppclass CRecordBatchReaderSourceNodeOptions
"arrow::acero::RecordBatchReaderSourceNodeOptions"(CExecNodeOptions):
+ CRecordBatchReaderSourceNodeOptions(shared_ptr[CRecordBatchReader]
reader)
+
cdef cppclass CSinkNodeOptions
"arrow::acero::SinkNodeOptions"(CExecNodeOptions):
pass
diff --git a/python/pyarrow/tests/test_acero.py
b/python/pyarrow/tests/test_acero.py
index e789b3bc7f..6e471d6119 100644
--- a/python/pyarrow/tests/test_acero.py
+++ b/python/pyarrow/tests/test_acero.py
@@ -25,6 +25,7 @@ try:
from pyarrow.acero import (
Declaration,
TableSourceNodeOptions,
+ RecordBatchReaderSourceNodeOptions,
FilterNodeOptions,
ProjectNodeOptions,
AggregateNodeOptions,
@@ -99,6 +100,96 @@ def test_table_source():
_ = decl.to_table()
+def test_record_batch_reader_source():
+ table = pa.table({'a': [1, 2, 3], 'b': [4, 5, 6]})
+ reader = pa.RecordBatchReader.from_batches(
+ table.schema, table.to_batches(max_chunksize=2)
+ )
+ decl = Declaration(
+ "record_batch_reader_source",
RecordBatchReaderSourceNodeOptions(reader)
+ )
+ result = decl.to_table()
+ assert result.equals(table)
+
+ # a reader can only be consumed once
+ decl = Declaration(
+ "record_batch_reader_source",
RecordBatchReaderSourceNodeOptions(reader)
+ )
+ result = decl.to_table()
+ assert result.num_rows == 0
+
+ with pytest.raises(TypeError):
+ RecordBatchReaderSourceNodeOptions(table)
+
+ with pytest.raises(TypeError):
+ RecordBatchReaderSourceNodeOptions(None)
+
+
+def test_record_batch_reader_source_lazy_generator():
+ # the reader can be backed by a Python generator, which is only
+ # consumed (from an I/O thread) while the plan executes
+ table = pa.table({'a': list(range(10)), 'b': list(range(10, 20))})
+ batches = table.to_batches(max_chunksize=2)
+ consumed = []
+
+ def gen():
+ for i, batch in enumerate(batches):
+ consumed.append(i)
+ yield batch
+
+ reader = pa.RecordBatchReader.from_batches(table.schema, gen())
+ decl = Declaration.from_sequence([
+ Declaration(
+ "record_batch_reader_source",
RecordBatchReaderSourceNodeOptions(reader)
+ ),
+ Declaration("filter", options=FilterNodeOptions(field("a") >= 5)),
+ ])
+ assert consumed == []
+ result = decl.to_table()
+ assert consumed == list(range(len(batches)))
+ assert result.sort_by("a").equals(table.slice(5))
+
+
+def test_record_batch_reader_source_generator_error():
+ # an error raised by the generator propagates to the plan execution
+ schema = pa.schema([("a", pa.int64())])
+
+ def gen():
+ yield pa.record_batch([pa.array([1, 2, 3])], schema=schema)
+ raise ValueError("error in generator")
+
+ reader = pa.RecordBatchReader.from_batches(schema, gen())
+ decl = Declaration(
+ "record_batch_reader_source",
RecordBatchReaderSourceNodeOptions(reader)
+ )
+ with pytest.raises(ValueError, match="error in generator"):
+ _ = decl.to_table()
+
+
+def test_record_batch_reader_source_hash_join_probe():
+ # streaming reader as the probe side of a hash join
+ left = pa.table({'key': [1, 2, 3, 4], 'a': ["a", "b", "c", "d"]})
+ right = pa.table({'key': [2, 3, 4, 5], 'b': ["p", "q", "r", "s"]})
+ reader = pa.RecordBatchReader.from_batches(
+ left.schema, left.to_batches(max_chunksize=1)
+ )
+ left_source = Declaration(
+ "record_batch_reader_source",
RecordBatchReaderSourceNodeOptions(reader)
+ )
+ right_source = Declaration("table_source", TableSourceNodeOptions(right))
+ join_opts = HashJoinNodeOptions(
+ "inner", left_keys="key", right_keys="key",
+ left_output=["key", "a"], right_output=["b"]
+ )
+ joined = Declaration("hashjoin", options=join_opts,
+ inputs=[left_source, right_source])
+ result = joined.to_table()
+ expected = pa.table({
+ 'key': [2, 3, 4], 'a': ["b", "c", "d"], 'b': ["p", "q", "r"]
+ })
+ assert result.sort_by("key").equals(expected)
+
+
def test_filter(table_source):
# referencing unknown field
decl = Declaration.from_sequence([
