From cea9f0180b7de1aaca88962e4ecaeba365e3baa6 Mon Sep 17 00:00:00 2001 From: Kevin Jacobs Date: Fri, 27 Mar 2026 03:07:57 -0400 Subject: [PATCH] Add batch_size parameter to read_parquet The default ParquetRecordBatchReaderBuilder batch_size of 1024 rows creates thousands of tiny batches for typical files. This adds a batch_size keyword argument to read_parquet, passed through to ParquetRecordBatchReaderBuilder::with_batch_size. On a 2M-row single-column file, read_all() drops from 4.3ms to 0.24ms with batch_size=65536. On a 5-column file, 46ms to 29ms. Signed-off-by: Kevin Jacobs --- arro3-io/python/arro3/io/_parquet.pyi | 11 ++++++++++- arro3-io/src/parquet.rs | 12 ++++++++++-- tests/io/test_parquet.py | 22 ++++++++++++++++++++++ 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/arro3-io/python/arro3/io/_parquet.pyi b/arro3-io/python/arro3/io/_parquet.pyi index 6f2bc61..4dceb1e 100644 --- a/arro3-io/python/arro3/io/_parquet.pyi +++ b/arro3-io/python/arro3/io/_parquet.pyi @@ -31,12 +31,21 @@ ParquetEncoding = Literal[ ] """Allowed Parquet encodings.""" -def read_parquet(file: IO[bytes] | Path | str) -> core.RecordBatchReader: +def read_parquet( + file: IO[bytes] | Path | str, + *, + batch_size: int | None = None, +) -> core.RecordBatchReader: """Read a Parquet file to an Arrow RecordBatchReader Args: file: The input Parquet file path or buffer. + Keyword Args: + batch_size: The number of rows per batch in the returned reader. + Defaults to 1024 if not specified. Larger values reduce per-batch + overhead and can significantly improve read performance. + Returns: The loaded Arrow data. """ diff --git a/arro3-io/src/parquet.rs b/arro3-io/src/parquet.rs index 13bd1fc..2816a31 100644 --- a/arro3-io/src/parquet.rs +++ b/arro3-io/src/parquet.rs @@ -24,8 +24,16 @@ use crate::error::Arro3IoResult; use crate::utils::{FileReader, FileWriter}; #[pyfunction] -pub fn read_parquet(file: FileReader) -> PyArrowResult { - let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); +#[pyo3(signature = (file, *, batch_size=None))] +pub fn read_parquet( + file: FileReader, + batch_size: Option, +) -> PyArrowResult { + let mut builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap(); + + if let Some(batch_size) = batch_size { + builder = builder.with_batch_size(batch_size); + } let metadata = builder.schema().metadata().clone(); let reader = builder.build().unwrap(); diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py index 3daa824..71ebe5e 100644 --- a/tests/io/test_parquet.py +++ b/tests/io/test_parquet.py @@ -51,3 +51,25 @@ def test_copy_parquet_kv_metadata(): reader = read_parquet(pq_path) assert reader.schema.metadata[b"hello"] == b"world" + + +def test_read_parquet_batch_size(): + table = pa.table({"a": list(range(10_000)), "b": ["x"] * 10_000}) + with TemporaryDirectory() as tmp_path: + tmp_path = Path(tmp_path) + pq_path = tmp_path / "test.parquet" + write_parquet(table, pq_path) + + # Default batch_size (1024) produces many batches + batches_default = list(read_parquet(pq_path)) + assert len(batches_default) > 1 + + # Large batch_size produces fewer batches + batches_large = list(read_parquet(pq_path, batch_size=10_000)) + assert len(batches_large) == 1 + assert batches_large[0].num_rows == 10_000 + + # Data is identical regardless of batch_size + table_default = pa.table(read_parquet(pq_path).read_all()) + table_batched = pa.table(read_parquet(pq_path, batch_size=5_000).read_all()) + assert table_default == table_batched