GH-50801: [Python] Expose the record_batch_reader_source Acero node (RecordBatchReaderSourceNodeOptions)

Project: apache/arrow | Language: C++ | Concepts: api-design, testing-strategy, data-structures | Difficulty: mid
Link: https://github.com/apache/arrow/pull/50802

The problem

The Acero execution engine in Apache Arrow did not expose the record_batch_reader_source node to Python. This meant users could not construct execution plans in Python that stream data from a RecordBatchReader, limiting flexibility for streaming and generator-backed data sources, especially for use cases like hash joins where streaming the probe side is more efficient.

Why it's non-trivial

A naive approach might try to wrap a RecordBatchReader directly, but this would not integrate with the Acero execution plan API, nor allow the use of this node as a first-class source in plan declarations. Proper exposure requires careful API design, C++/Python bindings, and robust testing to ensure correctness and Pythonic error handling.

The approach

1. Expose the node in the Python API documentation and import path:

+   RecordBatchReaderSourceNodeOptions

This line adds the new options class to the public API docs, making it discoverable.

+        RecordBatchReaderSourceNodeOptions,

This line imports the new class in the main Python module, so users can access it directly.

2. Cython bindings for the C++ class:

+cdef class _RecordBatchReaderSourceNodeOptions(ExecNodeOptions):
+
+    def _set_options(self, RecordBatchReader reader):
+        self.wrapped.reset(
+            new CRecordBatchReaderSourceNodeOptions(reader.reader)
+        )

This defines a Cython class that wraps the C++ RecordBatchReaderSourceNodeOptions, allowing a Python RecordBatchReader to be passed in and properly wrapped for use in execution plans.

3. Python-facing options class:

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

This exposes the node to Python users, with docstrings and parameter checks. It ensures only a RecordBatchReader can be passed.

4. Cython declaration for the C++ class:

+    cdef cppclass CRecordBatchReaderSourceNodeOptions "arrow::acero::RecordBatchReaderSourceNodeOptions"(CExecNodeOptions):
+        CRecordBatchReaderSourceNodeOptions(shared_ptr[CRecordBatchReader] reader)

This makes the C++ constructor available to Cython, enabling the Python bindings.

5. Comprehensive tests:

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

This test checks basic usage, single-consumption semantics, and type errors.

+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
+    ...
+    def gen():
+        for i, batch in enumerate(batches):
+            consumed.append(i)
+            yield batch
+    ...
+    assert consumed == []
+    result = decl.to_table()
+    assert consumed == list(range(len(batches)))
+    assert result.sort_by("a").equals(table.slice(5))

This test ensures lazy consumption of generator-backed readers.

+def test_record_batch_reader_source_generator_error():
+    # an error raised by the generator propagates to the plan execution
+    ...
+    def gen():
+        yield pa.record_batch([pa.array([1, 2, 3])], schema=schema)
+        raise ValueError("error in generator")
+    ...
+    with pytest.raises(ValueError, match="error in generator"):
+        _ = decl.to_table()

This verifies that generator errors propagate correctly.

+def test_record_batch_reader_source_hash_join_probe():
+    # streaming reader as the probe side of a hash join
+    ...
+    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)

This test demonstrates the main motivating use case: streaming the probe side of a hash join.

From the review thread

No notable review discussion on this PR.

Takeaways

• When exposing new functionality to a language binding, provide thorough tests that cover both typical and edge-case usage, including error propagation.
• API design for streaming data should consider consumption semantics (e.g., single-use readers) and integrate smoothly with existing plan/declaration mechanisms.
• Properly document and expose new options in both code and user-facing documentation to ensure discoverability and correct usage.