Skip to content

fix(io): stream RecordBatchReader writes through a rolling ParquetWriter#3661

Open
Sushankthatipally wants to merge 1 commit into
apache:mainfrom
Sushankthatipally:streaming-write-rolling-parquet-writer
Open

fix(io): stream RecordBatchReader writes through a rolling ParquetWriter#3661
Sushankthatipally wants to merge 1 commit into
apache:mainfrom
Sushankthatipally:streaming-write-rolling-parquet-writer

Conversation

@Sushankthatipally

Copy link
Copy Markdown

Rationale for this change

PR #3335 added pa.RecordBatchReader support to Table.append/Table.overwrite using a buffered bin-pack approach (bin_pack_record_batches). That implementation had two caveats that were already called out in its own docstring:

  1. Peak memory scaled with N_workers x write.target-file-size-bytes (roughly 4 GiB at defaults) instead of staying constant.
  2. write.target-file-size-bytes was measured in uncompressed in-memory Arrow bytes, not compressed on-disk Parquet bytes, so the resulting files ended up 3-10x smaller than the property suggested.

This PR replaces the buffered approach with a rolling pq.ParquetWriter driven by OutputStream.tell(), which was added in #2998 specifically for this purpose:

with output_file.create(overwrite=True) as fos:
    with pq.ParquetWriter(fos, schema=..., ...) as writer:
        writer.write_batch(first_batch)
        while fos.tell() < target_file_size:
            batch = next(batches)
            writer.write_batch(batch)

Both caveats go away. tell() reports actual compressed on-disk bytes, so write.target-file-size-bytes now matches the spec's intent (and the Java/Spark/Flink writers). Peak memory is bounded by one input batch plus the Parquet writer's internal page buffer, regardless of dataset size or how many files end up getting produced.

Row groups are capped at write.parquet.row-group-limit the same way the existing materialised pa.Table write path handles it: a batch bigger than the cap gets split into multiple row groups, a smaller batch becomes one row group. Callers control the lower bound through their input batch size, this just enforces the upper bound.

No public API change. tbl.append(reader) / tbl.overwrite(reader) work exactly the same from the caller's side, streaming writes are still unpartitioned-table only (tracked separately in #2152).

Are these changes tested?

Added two new tests in tests/catalog/test_catalog_behaviors.py:

  • test_append_record_batch_reader_row_group_limit_is_cap: a single batch bigger than write.parquet.row-group-limit gets split into the expected number of row groups, each exactly at the cap.
  • test_append_record_batch_reader_target_file_size_is_on_disk_bytes: streams several batches with a small write.target-file-size-bytes, checks multiple files get rolled, and asserts each rolled file lands close to the target (catching the old behaviour where files came out 3-10x smaller than the target).

Removed the bin_pack_record_batches unit tests since that function no longer exists. All existing streaming-write tests in test_catalog_behaviors.py still pass against both the in-memory and SQL catalog fixtures.

Are there any user-facing changes?

Yes, write.target-file-size-bytes now actually controls on-disk file size for streaming writes the way its docstring already claimed it should. Updated the docstrings on append/overwrite and the mkdocs streaming-writes doc page to describe the new rolling-writer behaviour instead of the old bin-packing one.

Closes apache#3388

PR apache#3335 added pa.RecordBatchReader support to Table.append/overwrite
using a buffered bin-pack approach (bin_pack_record_batches). That had
two problems called out in its own docstring: peak memory scaled with
worker count times target_file_size instead of staying constant, and
write.target-file-size-bytes was measured in uncompressed in-memory
Arrow bytes rather than actual on-disk Parquet bytes, so files ended
up 3-10x smaller than the property suggested.

This replaces that path with a rolling pq.ParquetWriter driven by
OutputStream.tell(), which was added in apache#2998 for exactly this reason.
Batches are written directly, a new file rolls once the on-disk byte
count crosses target_file_size, and row groups are capped at
write.parquet.row-group-limit the same way the materialised pa.Table
path already handles it. Peak memory is now bounded by one input batch
plus the writer's internal page buffer, independent of dataset size or
how many files get produced.

No public API change, tbl.append(reader) and tbl.overwrite(reader)
work the same way. Added tests covering the row group cap and the
on-disk byte accuracy of the rollover, and removed the old
bin_pack_record_batches tests since that function is gone.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant