GH-39788: [Python] Validate max_chunksize in Table.to_batches (#39796)

jorisvandenbossche · web-flow · commit 129a5291a26e · 2024-02-02T11:26:57.000+01:00
### Rationale for this change Validating the keyword to be strictly positive, to avoid an infinite loop. * Closes: #39788 Authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com> Signed-off-by: Raúl Cumplido <raulcumplido@gmail.com>
diff --git a/python/pyarrow/table.pxi b/python/pyarrow/table.pxi
@@ -4172,6 +4172,8 @@ cdef class Table(_Tabular):
         reader.reset(new TableBatchReader(deref(self.table)))
 
         if max_chunksize is not None:
+            if not max_chunksize > 0:
+                raise ValueError("'max_chunksize' should be strictly positive")
             c_max_chunksize = max_chunksize
             reader.get().set_chunksize(c_max_chunksize)
 
diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py
@@ -1089,6 +1089,9 @@ def test_table_to_batches():
     table_from_iter = pa.Table.from_batches(iter([batch1, batch2, batch1]))
     assert table.equals(table_from_iter)
 
+    with pytest.raises(ValueError):
+        table.to_batches(max_chunksize=0)
+
 
 def test_table_basics():
     data = [
