GH-49305: [Python] Expose RecordBatchFileReader.count_rows#50646
Draft
GujaLomsadze wants to merge 1 commit into
Draft
GH-49305: [Python] Expose RecordBatchFileReader.count_rows#50646GujaLomsadze wants to merge 1 commit into
GujaLomsadze wants to merge 1 commit into
Conversation
RecordBatchFileReader::CountRows was already available in Arrow C++ but not bound in Python, so the only way to get the total number of rows of an IPC file was to read every record batch and sum num_rows. Add a count_rows() method that reads only the metadata of each record batch without deserializing the batches themselves. The method is named after the C++ API and matches the existing count_rows() methods on Dataset, Scanner and Fragment.
|
|
|
|
1 similar comment
|
|
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
Resolves 49305
RecordBatchFileReader::CountRowshas existed in Arrow C++ (cpp/src/arrow/ipc/reader.h) butwas never bound in Python, so the only way to get the total number of rows of an IPC file was:
This findings are done in original opened Issue #49305
That deserializes every record batch just to read its length, which is wasteful and becomes
expensive on remote filesystems.
What changes are included in this PR?
Adds
RecordBatchFileReader.count_rows():Three changes:
python/pyarrow/includes/libarrow.pxd: declareCResult[int64_t] CountRows()onCRecordBatchFileReader, which was the missing piece.python/pyarrow/ipc.pxi: addcount_rows()to_RecordBatchFileReader, released GIL aroundthe call, with the same closed reader guard used by the existing
statspropertypython/pyarrow/tests/test_ipc.py: tests.On the naming:
count_rows()follows the C++ method and is consistent with the existingcount_rows()onDataset,ScannerandFragment.To be precise about the benefit, since the issue describes it as reading the count from the
metadata: the C++ implementation still walks every block, but reads only each record batch's
flatbuffer message header to pick up its length, and never touches the data buffers. So this is
a reduction in bytes read rather than in the number of reads, and the gain shows up on remote
filesystems and on files with large batches rather than in a local in memory benchmark.
This is only added to the file reader. The stream reader has no footer and cannot count rows
without consuming the stream.
Are these changes tested?
Yes, two tests in
python/pyarrow/tests/test_ipc.py:test_file_count_rows: count matches the sum of the written batch lengths, and counting doesnot consume the reader (count,
read_all(), count again).test_file_count_rows_no_batches: a file with a schema but no batches counts 0.Locally
test_ipc.pypasses (72 tests) andtest_feather.pypasses (83 passed, 8 skipped,1 xfailed), the latter because the feather reader sits on the same file reader. The docstring
example was run and produces the output shown.
Are there any user-facing changes?
Yes, a new public method
RecordBatchFileReader.count_rows(). No existing behaviour changes.AI usage disclosure
I used Claude Code to locate the unbound C++ method and the place where the declaration was
missing, and to draft the binding, the docstring and the tests. I reviewed the result, rebuilt
PyArrow locally, ran the test suites quoted above, and checked the C++ implementation of
CountRowsmyself to confirm what it actually does before describing the benefit here.