owenowenisme opened a new issue, #50801:
URL: https://github.com/apache/arrow/issues/50801
### Describe the enhancement requested
### Description
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`:
```python
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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
