Skip to content

[Python] Expose the record_batch_reader_source Acero node (RecordBatchReaderSourceNodeOptions) #50801

Description

@owenowenisme

Describe the enhancement requested

Acero has a record_batch_reader_source node in C++ since #15183 (added for the R bindings, which use it as ExecNode_SourceNode). Python currently has no way to use it: pyarrow.acero only exposes table_source (fully-materialized input) and the dataset scan node.

I'd like to add the corresponding options class to pyarrow.acero:

reader = pa.RecordBatchReader.from_batches(schema, batch_generator())
source = Declaration("record_batch_reader_source",
                     RecordBatchReaderSourceNodeOptions(reader))

This is a thin Cython binding (no C++ changes): declare arrow::acero::RecordBatchReaderSourceNodeOptions in libarrow_acero.pxd, wrap it in _acero.pyx following the TableSourceNodeOptions pattern, and re-export it from pyarrow.acero.

Motivation

There is currently no way from Python to run an Acero plan over in-memory data
without pinning all of the input for the plan's lifetime:

  • table_source requires the whole input as a single Table up front.
  • Declaration("scan", ScanNodeOptions(pyarrow.dataset.dataset(list_of_tables))) feeds batches into the plan incrementally, but the dataset keeps references to every fragment until the scan finishes, so input memory is never reclaimed while the plan runs.

My use case is a hash join where the probe side arrives as many independent Table chunks (the reduce step of a distributed shuffle, build side as table_source, probe side streamed through record_batch_reader_source, output consumed incrementally via Declaration.to_reader(). The hash table is built once, and neither the probe input nor the join output is ever fully resident.

Benchmark: inner join of a 481 MiB probe side (6M rows, arriving as 12 chunks) against a 154 MiB build side (3M rows), measuring peak process memory (USS) above baseline. All four strategies produce identical results:

how the probe side is fed peak memory wall hash builds
pa.concat_tables(chunks).join(build) +1032 MiB 1.1 s 1
chunk.join(build) per chunk, dropping each chunk +466 MiB 12.8 s 12
scan node over pyarrow.dataset.dataset(chunks) +807 MiB 1.4 s 1
record_batch_reader_source + releasing generator +380 MiB 1.4 s 1

The first row pays for materializing the full join output; the second avoids that but rebuilds the hash table for every chunk (Table.join runs a complete new plan per call); the third builds once and streams the output but cannot release any
probe input. The new node is the only combination of a single hash build, streaming output, and progressive release of the input.

Component(s)

Python

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions