Priya2996 opened a new pull request, #54840: URL: https://github.com/apache/arrow/pull/54840
### Rationale for this change The asynchronous Parquet record batch generator decoded an entire row group into memory before slicing it into record batches. For files containing large row groups, this caused high peak memory usage even when a smaller `batch_size` was configured. ### What changes are included in this PR? - Reuse the existing RecordBatchReader logic to decode row groups incrementally. - Yield batches as they are decoded instead of materializing the complete row group first. - Preserve asynchronous execution and row-group readahead. - Avoid deadlocks when the same executor is configured for CPU and I/O work by disabling parallel column decoding in that case. - Keep the existing non-pre-buffered fallback unchanged. ### Are these changes tested? Yes. Tests verify that: - Row groups are returned as multiple batches of the configured size. - Reassembling the batches produces the original table. - Empty-column projections preserve row counts and batch sizes. - A single-thread CPU executor completes successfully. - A single executor shared by CPU and I/O work does not deadlock. - Dataset scans with multiple files, row groups, and readahead complete with a single I/O thread. ### Are there any user-facing changes? There are no public API changes. For pre-buffered Parquet dataset scans, the existing batch_size option now also limits the number of rows decoded per batch instead of only controlling how an already-decoded row group is sliced. This reduces peak memory usage for large row groups. ### Was AI used for this PR? AI was used to help trace the relevant call chain, review concurrency behavior, suggest tests, and draft portions of the implementation and PR description. The changes were reviewed before submission. ### Benchmark Workload: one 262,144-row group with 128 `int64` columns, read in 32,768-row batches with pre-buffering enabled. Results are medians of five runs using real time. | Implementation | Rows/s | Peak reader memory | |---|---:|---:| | Main eager path | 139.3M | 268.4 MB | | Incremental, threaded | 53.6M | 33.8 MB | | Incremental, serial | 39.2M | 34.3 MB | The threaded incremental path reduces peak reader memory by about 87% compared with main and is about 37% faster than the serial alternative (Using Google Benchmark wall-clock time). A 1,481,181-row, 650-column Parquet file with one row group was scanned in a fresh process for each case. | Path | Batch size | Arrow peak | RSS peak | |---|---:|---:|---:| | Eager Dataset decode | 65,536 | 8.290 GiB | 5.420 GiB | | Incremental Dataset decode | 65,536 | 0.877 GiB | 1.144 GiB | | Eager Dataset decode | 131,072 | 8.290 GiB | 6.834 GiB | | Incremental Dataset decode | 131,072 | 1.739 GiB | 1.873 GiB | The incremental path reduced Arrow peak memory by 89.4% at a 65,536-row batch size and by 79.0% at 131,072 rows. The eager path's peak was unchanged across batch sizes, while incremental memory scaled with the configured batch size. ### Associated Arrow Issue Links https://github.com/apache/arrow/issues/36778 https://github.com/apache/arrow/issues/20473 #### PyIceberg PyIceberg’s to_arrow_batch_reader ultimately uses PyArrow’s Dataset scanner and Parquet record batch generator. Consequently, it reaches the whole-row-group decoding path addressed by this PR. For Parquet files with large row groups, this can appear as a memory leak or cause sustained memory growth because millions of rows are materialized before the first configured batch is returned. With this change, the existing batch size is applied during decoding. This bounds decoded batch memory and addresses the memory spikes observed when using PyIceberg’s `to_arrow_batch_reader`. We encountered this behavior while investigating high memory usage in PyIceberg’s to_arrow_batch_reader(). Tracing the call path into PyArrow showed that the asynchronous Parquet scanner decoded an entire row group before applying the configured batch_size. This Arrow-level behavior can amplify memory usage when Parquet files contain large row groups. The change in this PR makes decoding incremental for the pre-buffered scanner path. This will also help address concerns similar to ours for decoding large row groups using PyIceberg’s `to_arrow_batch_reader` **PR code and description written by:** - [x] Human - [x] AI **Reviewed before submission by:** - [x] Human - [x] AI - [ ] Not reviewed -- 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]
