Scanning a registered pyarrow.RecordBatchReader buffers most of the stream when the query consumes it slowly #691

What happens?

When a pyarrow.RecordBatchReader is registered and read by a query whose sink is slower than the reader's producer (for example CREATE TABLE ... AS SELECT of wide strings), most of the stream ends up held in Arrow memory at once. That memory sits outside memory_limit, so a query with memory_limit = '512MB' can hold gigabytes of input.

Registering the same reader's raw C stream (reader.__arrow_c_stream__()) instead keeps about one batch per scanning thread.

To Reproduce

The script needs only duckdb and pyarrow. It streams 1,000-row batches of 10 KB strings from a generator into CREATE TABLE t AS SELECT * FROM source, with memory_limit = '512MB' and threads = 4. While the statement runs it samples pyarrow.total_allocated_bytes(), which counts live Arrow allocations.

"""Registering a pyarrow RecordBatchReader buffers most of the stream when the query's sink is slower
than the producer; registering the reader's raw C stream does not.

Requires only duckdb and pyarrow. Usage: python repro.py [reader|capsule|scanner]
"""

import os
import sys
import tempfile
import threading
import time

import duckdb
import pyarrow as pa
import pyarrow.dataset as ds

ROWS_PER_BATCH = 1_000
BATCHES = int(os.environ.get("BATCHES", "100"))
VALUE_BYTES = 10_000  # 1,000 rows x 10 KB per batch: about 1 GB per 100 batches

schema = pa.schema([("id", pa.int64()), ("payload", pa.string())])
produced = 0


def batches():
    global produced
    for index in range(BATCHES):
        produced += 1
        start = index * ROWS_PER_BATCH
        yield pa.record_batch(
            [
                pa.array(range(start, start + ROWS_PER_BATCH), pa.int64()),
                pa.array(["x" * VALUE_BYTES] * ROWS_PER_BATCH),
            ],
            schema=schema,
        )


def peak_while(run):
    peak = 0
    done = threading.Event()

    def sample():
        nonlocal peak
        while not done.is_set():
            peak = max(peak, pa.total_allocated_bytes())
            time.sleep(0.002)

    sampler = threading.Thread(target=sample)
    sampler.start()
    try:
        run()
    finally:
        done.set()
        sampler.join()
    return peak


def main(mode):
    batch_mib = ROWS_PER_BATCH * VALUE_BYTES / 2**20
    print(f"duckdb {duckdb.__version__}, pyarrow {pa.__version__}, mode={mode}, "
          f"{BATCHES} batches of {batch_mib:.1f} MiB")
    reader = pa.RecordBatchReader.from_batches(schema, batches())
    if mode == "scanner":
        # The wrapper DuckDB applies, without DuckDB: read one batch every 20 ms.
        def run():
            for batch in ds.Scanner.from_batches(reader).to_reader():
                time.sleep(0.02)
    else:
        with tempfile.TemporaryDirectory() as spill:
            con = duckdb.connect(config={"memory_limit": "512MB", "threads": 4, "temp_directory": spill})
            con.register("source", reader.__arrow_c_stream__() if mode == "capsule" else reader)

            def run():
                con.execute("CREATE TABLE t AS SELECT * FROM source")

            peak = peak_while(run)
            con.close()
            report(peak, batch_mib)
            return
    report(peak_while(run), batch_mib)


def report(peak, batch_mib):
    peak_mib = peak / 2**20
    print(f"peak live Arrow allocations: {peak_mib:.0f} MiB (about {peak_mib / batch_mib:.0f} batches held), "
          f"batches produced: {produced}")


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "reader")

python repro.py reader    # register the RecordBatchReader
python repro.py capsule   # register reader.__arrow_c_stream__()
python repro.py scanner   # pyarrow alone: Scanner.from_batches(reader).to_reader(), read slowly

Output on the latest stable release and the latest nightly (macOS 26, arm64, Python 3.13, pyarrow 26.0.0), with BATCHES=300 (about 3 GB streamed):

duckdb 1.5.6, pyarrow 26.0.0, mode=reader, 300 batches of 9.5 MiB
peak live Arrow allocations: 688 MiB (about 72 batches held), batches produced: 300
duckdb 1.5.6, pyarrow 26.0.0, mode=capsule, 300 batches of 9.5 MiB
peak live Arrow allocations: 48 MiB (about 5 batches held), batches produced: 300
duckdb 2.0.0.dev2610151406, pyarrow 26.0.0, mode=reader, 300 batches of 9.5 MiB
peak live Arrow allocations: 573 MiB (about 60 batches held), batches produced: 300
duckdb 2.0.0.dev2610151406, pyarrow 26.0.0, mode=capsule, 300 batches of 9.5 MiB
peak live Arrow allocations: 48 MiB (about 5 batches held), batches produced: 300

Expected: Live Arrow memory stays near one batch per scanning thread, as it does with the capsule. The stream is pulled only as fast as the query consumes it.

Earlier measurements on DuckDB 1.5.4 and pyarrow 24.0.0 (Linux 6.6 aarch64, glibc, Python 3.13.15, mimalloc pool):

More detail:

• The held memory is live until the statement ends. Afterwards total_allocated_bytes() is 0 while the table still exists, so the table itself does not keep the buffers.
• A narrow sink barely triggers it. CREATE TABLE t AS SELECT id FROM source held about 250 MiB, and a plain count(*) held about 60 MiB (the count(*) figure is from a macOS run).
• With the capsule, the peak stays flat as the stream grows. At 600K rows of 11 KB it was 210 MiB, against 2.8 GB with the reader.

Where it comes from

• PythonTableArrowArrayStreamFactory::Produce (src/duckdb_py/arrow/arrow_array_stream.cpp, v1.5.4) treats a RecordBatchReader as a __arrow_c_stream__ object. When pyarrow.dataset is importable, it re-imports the stream and wraps it in pyarrow.dataset.Scanner.from_batches(reader, columns=..., filter=...).to_reader(), passing no read-ahead options.
• Scanner.from_batches(...).to_reader() reads its source far ahead of a slow consumer. In python repro.py scanner, 98 of 100 batches were live while the consumer was still reading.
• In a separate pure-pyarrow check, 64–75 of 75 batches had been produced after 11 were consumed. That held with batch_readahead=0 or 1, fragment_readahead=0 or 1, and use_threads=False, so those options do not bound it. This part may deserve its own Arrow issue.
• A raw PyCapsule takes the direct path in Produce. DuckDB consumes the C stream itself, one batch per scanning thread, and filters after the scan.

Workaround: Register reader.__arrow_c_stream__() instead of the reader. Projection and filter pushdown into pyarrow is lost; DuckDB applies them after the scan.

Possible fixes

• Consume RecordBatchReader and other stream objects directly, as the PyCapsule path does, when there is nothing to push down. Or make that the default for one-shot streams.
• When wrapping in a Scanner, bound its read-ahead. This may need an Arrow change if from_batches ignores batch_readahead for one-shot sources.

OS:

macOS 26 arm64 (latest runs); Linux 6.6 aarch64, glibc (1.5.4 runs)

DuckDB Package Version:

1.5.6, and 2.0.0.dev2610151406

Python Version:

3.13

Full Name:

Max Parke

Affiliation:

Scale AI

What is the latest build you tested with? If possible, we recommend testing with the latest nightly build.

I have tested with a nightly build

Did you include all relevant data sets for reproducing the issue?

Not applicable - the reproduction does not require a data set

Did you include all code required to reproduce the issue?

• Yes, I have

Did you include all relevant configuration to reproduce the issue?

• Yes, I have

🤖 — posted via Claude Code