diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index e57c6d0d92d..efc9602a3a8 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 6477579af21..991505a2b7f 100644 --- a/python/pyarrow/ipc.pxi +++ b/python/pyarrow/ipc.pxi @@ -1192,6 +1192,39 @@ 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 + + if not self.reader: + raise ValueError("Operation on closed reader") + + 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 6813ed77723..121617371c9 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()