Skip to content

Commit bf79398

Browse files
authored
GH-46811: [C++][Python] Fix crash on FileReaderImpl::GetRecordBatchReader (#46931)
### Rationale for this change Fixes #46811. ### What changes are included in this PR? - Added check to return Invalid when batch_size <= 0. ### Are these changes tested? Exercised in CI. ### Are there any user-facing changes? **This PR contains a "Critical Fix".** * GitHub Issue: #46811 Authored-by: Bryce Mecum <petridish@gmail.com> Signed-off-by: Bryce Mecum <petridish@gmail.com>
1 parent eda9942 commit bf79398

2 files changed

Lines changed: 14 additions & 0 deletions

File tree

python/pyarrow/parquet/core.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,6 +579,9 @@ def iter_batches(self, batch_size=65536, row_groups=None, columns=None,
579579
4 5 Brittle stars
580580
5 100 Centipede
581581
"""
582+
if batch_size <= 0:
583+
raise ValueError("batch_size must be greater than zero")
584+
582585
if row_groups is None:
583586
row_groups = range(0, self.metadata.num_row_groups)
584587
column_indices = self._get_column_indices(

python/pyarrow/tests/parquet/test_parquet_file.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -430,3 +430,14 @@ def test_fsspec_uri_raises_if_fsspec_is_not_available():
430430
"`fsspec` is required to handle `fsspec+<filesystem>://` and `hf://` URIs.")
431431
with pytest.raises(ImportError, match=msg):
432432
pq.read_table("fsspec+memory://example.parquet")
433+
434+
435+
def test_iter_batches_raises_batch_size_zero(tempdir):
436+
# See https://github.com/apache/arrow/issues/46811
437+
schema = pa.schema([])
438+
empty_table = pa.Table.from_batches([], schema=schema)
439+
parquet_file_path = tempdir / "empty_file.parquet"
440+
pq.write_table(empty_table, parquet_file_path)
441+
parquet_file = pq.ParquetFile(parquet_file_path)
442+
with pytest.raises(ValueError):
443+
parquet_file.iter_batches(batch_size=0)

0 commit comments

Comments
 (0)