feat: add batch_size parameter to read_parquet

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.
     """

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<Arro3RecordBatchReader> {
-    let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
+#[pyo3(signature = (file, *, batch_size=None))]
+pub fn read_parquet(
+    file: FileReader,
+    batch_size: Option<usize>,
+) -> PyArrowResult<Arro3RecordBatchReader> {
+    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();

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