This is an automated email from the ASF dual-hosted git repository.
pitrou pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 84027295a1 GH-49305: [Python] Expose RecordBatchFileReader.count_rows
(#50646)
84027295a1 is described below
commit 84027295a1f5dd3a100ae4f5bc41ad68a45ff388
Author: Guja <[email protected]>
AuthorDate: Thu Jul 30 11:59:47 2026 +0200
GH-49305: [Python] Expose RecordBatchFileReader.count_rows (#50646)
### Rationale for this change
Resolves [49305](https://github.com/apache/arrow/issues/49305)
`RecordBatchFileReader::CountRows` has existed in Arrow C++
(`cpp/src/arrow/ipc/reader.h`) but
was never bound in Python, so the only way to get the total number of rows
of an IPC file was:
```python
num_rows = sum(reader.get_batch(i).num_rows for i in
range(reader.num_record_batches))
```
### 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()`:
```python
with pa.ipc.open_file(source) as reader:
reader.count_rows()
```
Three changes:
* `python/pyarrow/includes/libarrow.pxd`: declare `CResult[int64_t]
CountRows()` on
`CRecordBatchFileReader`, which was the missing piece.
* `python/pyarrow/ipc.pxi`: add `count_rows()` to `_RecordBatchFileReader`,
released GIL around
the call, with the same closed reader guard used by the existing `stats`
property
* `python/pyarrow/tests/test_ipc.py`: tests.
On the naming: `count_rows()` follows the C++ method and is consistent with
the existing
`count_rows()` on `Dataset`, `Scanner` and `Fragment`.
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 does
not 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.py` passes (72 tests) and `test_feather.py` passes (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
`CountRows` myself to confirm what it actually does before describing the
benefit here.
* GitHub Issue: #49305
Authored-by: Guja <[email protected]>
Signed-off-by: Antoine Pitrou <[email protected]>
---
python/pyarrow/includes/libarrow.pxd | 2 ++
python/pyarrow/ipc.pxi | 30 ++++++++++++++++++++++++++++++
python/pyarrow/tests/test_ipc.py | 24 ++++++++++++++++++++++++
3 files changed, 56 insertions(+)
diff --git a/python/pyarrow/includes/libarrow.pxd
b/python/pyarrow/includes/libarrow.pxd
index e57c6d0d92..efc9602a3a 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -2011,6 +2011,8 @@ cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc"
nogil:
CResult[CRecordBatchWithMetadata]
ReadRecordBatchWithCustomMetadata(int i)
+ CResult[int64_t] CountRows()
+
CIpcReadStats stats()
shared_ptr[const CKeyValueMetadata] metadata()
diff --git a/python/pyarrow/ipc.pxi b/python/pyarrow/ipc.pxi
index 6477579af2..2c0d9591c1 100644
--- a/python/pyarrow/ipc.pxi
+++ b/python/pyarrow/ipc.pxi
@@ -1192,6 +1192,36 @@ cdef class _RecordBatchFileReader(_Weakrefable):
"""
return self.reader.get().num_record_batches()
+ def count_rows(self):
+ """
+ The total number of rows in the IPC file.
+
+ This reads the metadata of each record batch in the file, without
+ deserializing the record batches themselves.
+
+ Returns
+ -------
+ count : int
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> schema = pa.schema([('a', pa.int64())])
+ >>> sink = pa.BufferOutputStream()
+ >>> with pa.ipc.new_file(sink, schema) as writer:
+ ... for i in range(3):
+ ... writer.write_batch(pa.record_batch([[1, 2]],
schema=schema))
+ >>> with pa.ipc.open_file(sink.getvalue()) as reader:
+ ... reader.count_rows()
+ 6
+ """
+ cdef int64_t nrows
+
+ with nogil:
+ nrows = GetResultValue(self.reader.get().CountRows())
+
+ return nrows
+
def get_batch(self, int i):
"""
Read the record batch with the given index.
diff --git a/python/pyarrow/tests/test_ipc.py b/python/pyarrow/tests/test_ipc.py
index 6813ed7772..121617371c 100644
--- a/python/pyarrow/tests/test_ipc.py
+++ b/python/pyarrow/tests/test_ipc.py
@@ -189,6 +189,30 @@ def test_file_read_all(sink_factory):
assert result.equals(expected)
+def test_file_count_rows(file_fixture):
+ batches = file_fixture.write_batches()
+ file_contents = pa.BufferReader(file_fixture.get_source())
+
+ reader = pa.ipc.open_file(file_contents)
+
+ expected = sum(batch.num_rows for batch in batches)
+ assert reader.count_rows() == expected
+ # counting the rows does not consume the reader
+ assert reader.read_all().num_rows == expected
+ assert reader.count_rows() == expected
+
+
+def test_file_count_rows_no_batches():
+ schema = pa.schema([('a', pa.int64())])
+ sink = pa.BufferOutputStream()
+ with pa.ipc.new_file(sink, schema):
+ pass
+
+ reader = pa.ipc.open_file(sink.getvalue())
+ assert reader.num_record_batches == 0
+ assert reader.count_rows() == 0
+
+
def test_open_file_from_buffer(file_fixture):
# ARROW-2859; APIs accept the buffer protocol
file_fixture.write_batches()
