fix(io): stream RecordBatchReader writes through a rolling ParquetWriter#3661
Open
Sushankthatipally wants to merge 1 commit into
Open
fix(io): stream RecordBatchReader writes through a rolling ParquetWriter#3661Sushankthatipally wants to merge 1 commit into
Sushankthatipally wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rationale for this change
PR #3335 added
pa.RecordBatchReadersupport toTable.append/Table.overwriteusing a buffered bin-pack approach (bin_pack_record_batches). That implementation had two caveats that were already called out in its own docstring:N_workers x write.target-file-size-bytes(roughly 4 GiB at defaults) instead of staying constant.write.target-file-size-byteswas 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.ParquetWriterdriven byOutputStream.tell(), which was added in #2998 specifically for this purpose:Both caveats go away.
tell()reports actual compressed on-disk bytes, sowrite.target-file-size-bytesnow 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-limitthe same way the existing materialisedpa.Tablewrite 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 thanwrite.parquet.row-group-limitgets 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 smallwrite.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_batchesunit tests since that function no longer exists. All existing streaming-write tests intest_catalog_behaviors.pystill pass against both the in-memory and SQL catalog fixtures.Are there any user-facing changes?
Yes,
write.target-file-size-bytesnow actually controls on-disk file size for streaming writes the way its docstring already claimed it should. Updated the docstrings onappend/overwriteand themkdocsstreaming-writes doc page to describe the new rolling-writer behaviour instead of the old bin-packing one.