You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by wj...@apache.org on 2023/03/06 23:10:00 UTC
[arrow] branch main updated: GH-34410: [Python] Allow chunk sizes larger than the default to be used (#34435)
This is an automated email from the ASF dual-hosted git repository.
wjones127 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new fc950192ae GH-34410: [Python] Allow chunk sizes larger than the default to be used (#34435)
fc950192ae is described below
commit fc950192ae7b801a76adc3b3de410dc22da83fd4
Author: Weston Pace <we...@gmail.com>
AuthorDate: Mon Mar 6 15:09:51 2023 -0800
GH-34410: [Python] Allow chunk sizes larger than the default to be used (#34435)
### Rationale for this change
We changed the default chunk size from 64Mi rows to 1Mi rows. However, it turns out that this property was being treated not just as the default but also as the absolute max. So it was no longer possible to specify chunk sizes larger than 1Mi rows. This change separates those two things and restores the max to 64Mi rows.
### What changes are included in this PR?
Pyarrow will now set the `ParquetWriterProperties::max_row_group_length` to 64Mi when constructing a parquet writer.
### Are these changes tested?
Yes. Unit tests are added.
### Are there any user-facing changes?
No. The previous change #34281 changed two defaults (absolute max and default). This PR restores the absolute max back to what it was before. So it is removing a user-facing change.
* Closes: #34410
Lead-authored-by: Weston Pace <we...@gmail.com>
Co-authored-by: Will Jones <wi...@gmail.com>
Signed-off-by: Will Jones <wi...@gmail.com>
---
python/pyarrow/_parquet.pxd | 1 +
python/pyarrow/_parquet.pyx | 9 ++++++
python/pyarrow/parquet/core.py | 6 ++--
python/pyarrow/tests/parquet/common.py | 4 +++
.../pyarrow/tests/parquet/test_parquet_writer.py | 37 +++++++++++++++++++++-
5 files changed, 54 insertions(+), 3 deletions(-)
diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd
index 3996b5ec23..a3db7ab8d3 100644
--- a/python/pyarrow/_parquet.pxd
+++ b/python/pyarrow/_parquet.pxd
@@ -417,6 +417,7 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil:
Builder* encoding(ParquetEncoding encoding)
Builder* encoding(const c_string& path,
ParquetEncoding encoding)
+ Builder* max_row_group_length(int64_t size)
Builder* write_batch_size(int64_t batch_size)
Builder* dictionary_pagesize_limit(int64_t dictionary_pagesize_limit)
shared_ptr[WriterProperties] build()
diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx
index 04fd71d590..5dfd818a3d 100644
--- a/python/pyarrow/_parquet.pyx
+++ b/python/pyarrow/_parquet.pyx
@@ -1597,6 +1597,15 @@ cdef shared_ptr[WriterProperties] _create_writer_properties(
props.encryption(
(<FileEncryptionProperties>encryption_properties).unwrap())
+ # For backwards compatibility reasons we cap the maximum row group size
+ # at 64Mi rows. This could be changed in the future, though it would be
+ # a breaking change.
+ #
+ # The user can always specify a smaller row group size (and the default
+ # is smaller) when calling write_table. If the call to write_table uses
+ # a size larger than this then it will be latched to this value.
+ props.max_row_group_length(64*1024*1024)
+
properties = props.build()
return properties
diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py
index eb81b976c9..018902dde4 100644
--- a/python/pyarrow/parquet/core.py
+++ b/python/pyarrow/parquet/core.py
@@ -1059,7 +1059,8 @@ Examples
row_group_size : int, default None
Maximum number of rows in written row group. If None, the
row group size will be the minimum of the RecordBatch
- size and 1024 * 1024.
+ size and 1024 * 1024. If set larger than 64Mi then 64Mi
+ will be used instead.
"""
table = pa.Table.from_batches([batch], batch.schema)
self.write_table(table, row_group_size)
@@ -1074,7 +1075,8 @@ Examples
row_group_size : int, default None
Maximum number of rows in each written row group. If None,
the row group size will be the minimum of the Table size
- and 1024 * 1024.
+ and 1024 * 1024. If set larger than 64Mi then 64Mi will
+ be used instead.
"""
if self.schema_changed:
diff --git a/python/pyarrow/tests/parquet/common.py b/python/pyarrow/tests/parquet/common.py
index 09bc1ed7b3..f4e609c6ff 100644
--- a/python/pyarrow/tests/parquet/common.py
+++ b/python/pyarrow/tests/parquet/common.py
@@ -111,6 +111,10 @@ def _random_integers(size, dtype):
size=size).astype(dtype)
+def _range_integers(size, dtype):
+ return pa.array(np.arange(size, dtype=dtype))
+
+
def _test_dataframe(size=10000, seed=0):
import pandas as pd
diff --git a/python/pyarrow/tests/parquet/test_parquet_writer.py b/python/pyarrow/tests/parquet/test_parquet_writer.py
index 83cbac5936..6ae4307135 100644
--- a/python/pyarrow/tests/parquet/test_parquet_writer.py
+++ b/python/pyarrow/tests/parquet/test_parquet_writer.py
@@ -24,7 +24,8 @@ from pyarrow.tests.parquet.common import parametrize_legacy_dataset
try:
import pyarrow.parquet as pq
- from pyarrow.tests.parquet.common import _read_table, _test_dataframe
+ from pyarrow.tests.parquet.common import (_read_table, _test_dataframe,
+ _range_integers)
except ImportError:
pq = None
@@ -202,6 +203,40 @@ def test_parquet_writer_write_wrappers(tempdir, filesystem):
tm.assert_frame_equal(result, df)
+@pytest.mark.large_memory
+@pytest.mark.pandas
+def test_parquet_writer_chunk_size(tempdir):
+ default_chunk_size = 1024 * 1024
+ abs_max_chunk_size = 64 * 1024 * 1024
+
+ def check_chunk_size(data_size, chunk_size, expect_num_chunks):
+ table = pa.Table.from_arrays([
+ _range_integers(data_size, 'b')
+ ], names=['x'])
+ pq.write_table(table, tempdir / 'test.parquet', row_group_size=chunk_size)
+ metadata = pq.read_metadata(tempdir / 'test.parquet')
+ assert metadata.num_row_groups == expect_num_chunks
+ latched_chunk_size = min(chunk_size, abs_max_chunk_size)
+ # First chunks should be full size
+ for chunk_idx in range(expect_num_chunks - 1):
+ assert metadata.row_group(chunk_idx).num_rows == latched_chunk_size
+ # Last chunk may be smaller
+ remainder = data_size - (chunk_size * (expect_num_chunks - 1))
+ if remainder == 0:
+ assert metadata.row_group(
+ expect_num_chunks - 1).num_rows == latched_chunk_size
+ else:
+ assert metadata.row_group(expect_num_chunks - 1).num_rows == remainder
+
+ check_chunk_size(default_chunk_size * 2, default_chunk_size - 100, 3)
+ check_chunk_size(default_chunk_size * 2, default_chunk_size, 2)
+ check_chunk_size(default_chunk_size * 2, default_chunk_size + 100, 2)
+ check_chunk_size(default_chunk_size + 100, default_chunk_size + 100, 1)
+ # Even though the chunk size requested is large enough it will be capped
+ # by the absolute max chunk size
+ check_chunk_size(abs_max_chunk_size * 2, abs_max_chunk_size * 2, 2)
+
+
@pytest.mark.pandas
@pytest.mark.parametrize("filesystem", [
None,