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
Describe the enhancement requested
Acero has a
record_batch_reader_sourcenode in C++ since #15183 (added for the R bindings, which use it asExecNode_SourceNode). Python currently has no way to use it:pyarrow.aceroonly exposestable_source(fully-materialized input) and the datasetscannode.I'd like to add the corresponding options class to
pyarrow.acero:This is a thin Cython binding (no C++ changes): declare
arrow::acero::RecordBatchReaderSourceNodeOptionsinlibarrow_acero.pxd, wrap it in_acero.pyxfollowing theTableSourceNodeOptionspattern, and re-export it frompyarrow.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_sourcerequires the whole input as a singleTableup 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
Tablechunks (the reduce step of a distributed shuffle, build side astable_source, probe side streamed throughrecord_batch_reader_source, output consumed incrementally viaDeclaration.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:
pa.concat_tables(chunks).join(build)chunk.join(build)per chunk, dropping each chunkscannode overpyarrow.dataset.dataset(chunks)record_batch_reader_source+ releasing generatorThe first row pays for materializing the full join output; the second avoids that but rebuilds the hash table for every chunk (
Table.joinruns a complete new plan per call); the third builds once and streams the output but cannot release anyprobe input. The new node is the only combination of a single hash build, streaming output, and progressive release of the input.
Component(s)
Python