From 4569bf4bd2f45cca70a585ac2d5abe26f0d54d12 Mon Sep 17 00:00:00 2001 From: You-Cheng Lin Date: Tue, 4 Aug 2026 16:53:07 +0800 Subject: [PATCH] update Signed-off-by: You-Cheng Lin --- 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 b75fc330615e..363353bf9f6f 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 6e88fc7210da..e27509469ef5 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 e475e8db5c24..55f03a8bb75f 100644 --- a/python/pyarrow/acero.py +++ b/python/pyarrow/acero.py @@ -30,6 +30,7 @@ Declaration, ExecNodeOptions, TableSourceNodeOptions, + RecordBatchReaderSourceNodeOptions, FilterNodeOptions, ProjectNodeOptions, AggregateNodeOptions, diff --git a/python/pyarrow/includes/libarrow_acero.pxd b/python/pyarrow/includes/libarrow_acero.pxd index dc9babee190e..7e81a393682c 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 e789b3bc7fa2..6e471d61199f 100644 --- a/python/pyarrow/tests/test_acero.py +++ b/python/pyarrow/tests/test_acero.py @@ -25,6 +25,7 @@ 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([