From f59e1a6bca94b496d6a8bbfc044ed58ff9ffd375 Mon Sep 17 00:00:00 2001 From: Bekbaganbetov Abay Date: Tue, 21 Jul 2026 17:03:22 -0700 Subject: [PATCH] [Python][Parquet] Expose size statistics level in the Parquet writer Add a `write_size_statistics` keyword to `pyarrow.parquet.write_table`, `ParquetWriter`, and the dataset write path. It exposes the C++ `parquet::WriterProperties::Builder::set_size_statistics_level` setter introduced in GH-40592, accepting "none", "columnchunk", or "pageandcolumnchunk" (None keeps the Arrow C++ default of "pageandcolumnchunk"). Size statistics are written by default, and until now there was no way to disable or tune them from PyArrow. "none" turns them off entirely; page level statistics additionally require the page index to be enabled. --- python/pyarrow/_dataset_parquet.pyx | 2 + python/pyarrow/_parquet.pxd | 1 + python/pyarrow/_parquet.pyx | 23 ++++- python/pyarrow/includes/libparquet.pxd | 9 ++ python/pyarrow/parquet/core.py | 15 +++ python/pyarrow/tests/parquet/test_basic.py | 107 +++++++++++++++++++++ 6 files changed, 155 insertions(+), 2 deletions(-) diff --git a/python/pyarrow/_dataset_parquet.pyx b/python/pyarrow/_dataset_parquet.pyx index 534f7790923a..4812f567af37 100644 --- a/python/pyarrow/_dataset_parquet.pyx +++ b/python/pyarrow/_dataset_parquet.pyx @@ -660,6 +660,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): write_page_checksum=self._properties["write_page_checksum"], sorting_columns=self._properties["sorting_columns"], store_decimal_as_integer=self._properties["store_decimal_as_integer"], + write_size_statistics=self._properties["write_size_statistics"], ) def _set_arrow_properties(self): @@ -713,6 +714,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): write_page_checksum=False, sorting_columns=None, store_decimal_as_integer=False, + write_size_statistics=None, ) self._set_properties() diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd index 36fc2ccf2f33..939723419307 100644 --- a/python/pyarrow/_parquet.pxd +++ b/python/pyarrow/_parquet.pxd @@ -56,6 +56,7 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( write_page_checksum=*, sorting_columns=*, store_decimal_as_integer=*, + write_size_statistics=*, use_content_defined_chunking=*, bloom_filter_options=* ) except * diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index 2358a961ebd9..7367c7c62693 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -2066,6 +2066,7 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( write_page_checksum=False, sorting_columns=None, store_decimal_as_integer=False, + write_size_statistics=None, use_content_defined_chunking=False, bloom_filter_options=None) except *: @@ -2269,6 +2270,22 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( else: props.disable_write_page_index() + # size statistics + + if write_size_statistics is not None: + if write_size_statistics == "none": + props.set_size_statistics_level(SizeStatisticsLevel_None) + elif write_size_statistics == "columnchunk": + props.set_size_statistics_level(SizeStatisticsLevel_ColumnChunk) + elif write_size_statistics == "pageandcolumnchunk": + props.set_size_statistics_level( + SizeStatisticsLevel_PageAndColumnChunk) + else: + raise ValueError( + "Unsupported size statistics level: " + f"{write_size_statistics!r}. Expected one of 'none', " + "'columnchunk', 'pageandcolumnchunk'.") + properties = props.build() return properties @@ -2396,7 +2413,8 @@ cdef class ParquetWriter(_Weakrefable): store_decimal_as_integer=False, use_content_defined_chunking=False, write_time_adjusted_to_utc=False, - bloom_filter_options=None): + bloom_filter_options=None, + write_size_statistics=None): cdef: shared_ptr[WriterProperties] properties shared_ptr[ArrowWriterProperties] arrow_properties @@ -2433,7 +2451,8 @@ cdef class ParquetWriter(_Weakrefable): sorting_columns=sorting_columns, store_decimal_as_integer=store_decimal_as_integer, use_content_defined_chunking=use_content_defined_chunking, - bloom_filter_options=bloom_filter_options + bloom_filter_options=bloom_filter_options, + write_size_statistics=write_size_statistics ) arrow_properties = _create_arrow_writer_properties( use_deprecated_int96_timestamps=use_deprecated_int96_timestamps, diff --git a/python/pyarrow/includes/libparquet.pxd b/python/pyarrow/includes/libparquet.pxd index df353cc7805f..6d5125f5c844 100644 --- a/python/pyarrow/includes/libparquet.pxd +++ b/python/pyarrow/includes/libparquet.pxd @@ -516,6 +516,7 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil: Builder* disable_bloom_filter(const c_string& path) Builder* enable_bloom_filter(const c_string& path, BloomFilterOptions bloom_filter_options) + Builder* set_size_statistics_level(SizeStatisticsLevel level) shared_ptr[WriterProperties] build() cdef cppclass ArrowWriterProperties: @@ -599,6 +600,14 @@ cdef extern from "parquet/arrow/schema.h" namespace "parquet::arrow" nogil: cdef extern from "parquet/properties.h" namespace "parquet" nogil: + cdef enum SizeStatisticsLevel: + SizeStatisticsLevel_None \ + " parquet::SizeStatisticsLevel::None" + SizeStatisticsLevel_ColumnChunk \ + " parquet::SizeStatisticsLevel::ColumnChunk" + SizeStatisticsLevel_PageAndColumnChunk \ + " parquet::SizeStatisticsLevel::PageAndColumnChunk" + cdef enum ArrowWriterEngineVersion: V1 "parquet::ArrowWriterProperties::V1", V2 "parquet::ArrowWriterProperties::V2" diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index ff880fdcf52c..f109afe4027c 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -895,6 +895,17 @@ def _sanitize_table(table, new_schema, flavor): Whether to write page checksums in general for all columns. Page checksums enable detection of data corruption, which might occur during transmission or in the storage. +write_size_statistics : {"none", "columnchunk", "pageandcolumnchunk"}, default None + Control the level of size statistics written to the Parquet file. Size + statistics record the (unencoded) byte sizes and definition/repetition + level histograms of the data. If None, the Arrow C++ default is used + (currently "pageandcolumnchunk"). Use "none" to disable writing size + statistics entirely, "columnchunk" to write them only at the column-chunk + level, or "pageandcolumnchunk" to also write them into the page index. + Note that page-level size statistics are only written when the page index + is also enabled (see ``write_page_index``); with the default + ``write_page_index=False``, "pageandcolumnchunk" behaves like + "columnchunk". sorting_columns : Sequence of SortingColumn, default None Specify the sort order of the data being written. The writer does not sort the data nor does it verify that the data is sorted. The sort order is @@ -1083,6 +1094,7 @@ def __init__(self, where, schema, filesystem=None, store_decimal_as_integer=False, write_time_adjusted_to_utc=False, max_rows_per_page=None, + write_size_statistics=None, **options): if use_deprecated_int96_timestamps is None: # Use int96 timestamps for Spark @@ -1138,6 +1150,7 @@ def __init__(self, where, schema, filesystem=None, store_decimal_as_integer=store_decimal_as_integer, write_time_adjusted_to_utc=write_time_adjusted_to_utc, max_rows_per_page=max_rows_per_page, + write_size_statistics=write_size_statistics, **options) self.is_open = True @@ -2017,6 +2030,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_time_adjusted_to_utc=False, max_rows_per_page=None, bloom_filter_options=None, + write_size_statistics=None, **kwargs): # Implementor's note: when adding keywords here / updating defaults, also # update it in write_to_dataset and _dataset_parquet.pyx ParquetFileWriteOptions @@ -2051,6 +2065,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_time_adjusted_to_utc=write_time_adjusted_to_utc, max_rows_per_page=max_rows_per_page, bloom_filter_options=bloom_filter_options, + write_size_statistics=write_size_statistics, **kwargs) as writer: writer.write_table(table, row_group_size=row_group_size) except Exception: diff --git a/python/pyarrow/tests/parquet/test_basic.py b/python/pyarrow/tests/parquet/test_basic.py index 20e3f51bb677..f5c585c38f86 100644 --- a/python/pyarrow/tests/parquet/test_basic.py +++ b/python/pyarrow/tests/parquet/test_basic.py @@ -934,6 +934,113 @@ def test_thrift_size_limits(tempdir): assert got == table +_SIZE_STATISTICS_LEVELS = ["none", "columnchunk", "pageandcolumnchunk"] + + +def _size_statistics_table(nrows=5000): + # Nulls give non-empty definition-level histograms and variable-length + # strings give non-zero unencoded byte sizes, so the size statistics are + # non-trivial and their presence is observable in the file size. + return pa.table({ + 's': pa.array([("x" * 10 if i % 3 else None) for i in range(nrows)], + type=pa.string()), + 'i': pa.array(range(nrows), type=pa.int64()), + }) + + +def _write_parquet_bytes(table, **kwargs): + buf = io.BytesIO() + pq.write_table(table, buf, **kwargs) + return buf.getvalue() + + +@pytest.mark.parametrize("level", _SIZE_STATISTICS_LEVELS) +@pytest.mark.parametrize("use_writer", [False, True]) +def test_write_size_statistics_levels(tempdir, level, use_writer): + """Every level round-trips, via both write_table and ParquetWriter.""" + table = pa.table({'a': [1, 2, 3, 4], 'b': ['x', 'y', 'z', None]}) + path = tempdir / f'size_stats_{level}_{use_writer}.parquet' + if use_writer: + with pq.ParquetWriter(path, table.schema, + write_size_statistics=level) as writer: + writer.write_table(table) + else: + pq.write_table(table, path, write_size_statistics=level) + assert pq.read_table(path) == table + + +@pytest.mark.parametrize("level", _SIZE_STATISTICS_LEVELS) +def test_write_size_statistics_nested_and_null(tempdir, level): + """Levels work across nullable, empty, and nested/list columns.""" + table = pa.table({ + 'str': pa.array(['a', None, 'ccc', ''], pa.string()), + 'lst': pa.array([[1, 2], None, [], [3]], pa.list_(pa.int64())), + 'i32': pa.array([1, None, 3, 4], pa.int32()), + }) + path = tempdir / f'size_stats_nested_{level}.parquet' + pq.write_table(table, path, write_size_statistics=level) + assert pq.read_table(path) == table + + +@pytest.mark.parametrize( + "value", ["not-a-level", "None", "COLUMNCHUNK", "page_and_column_chunk", ""]) +def test_write_size_statistics_invalid(value): + """Unknown or mis-cased size statistics levels are rejected.""" + table = pa.table({'a': [1, 2, 3, 4]}) + with pytest.raises(ValueError, match="size statistics level"): + pq.write_table(table, io.BytesIO(), write_size_statistics=value) + + +def test_write_size_statistics_dataset(tempdir): + """The option also flows through the dataset write path.""" + table = pa.table({'a': [1, 2, 3, 4], 'b': ['x', 'y', 'z', None]}) + pq.write_to_dataset(table, str(tempdir / 'ds'), + write_size_statistics='none') + assert pq.read_table(str(tempdir / 'ds')).sort_by('a') == table + + +def test_write_size_statistics_reduces_file_size(): + """Disabling size statistics measurably shrinks the file, and the + column-chunk level omits the page-level statistics that the + page-and-column-chunk level adds (only when the page index is written).""" + table = _size_statistics_table() + + # Column-chunk size statistics live in the (uncompressed) footer, so + # "none" is strictly smaller than "columnchunk" regardless of page index. + none = len(_write_parquet_bytes(table, write_size_statistics="none")) + cc = len(_write_parquet_bytes(table, write_size_statistics="columnchunk")) + assert none < cc + + # Page-level size statistics are only emitted when the page index is on, + # so "pageandcolumnchunk" adds bytes over "columnchunk" only in that case. + cc_idx = len(_write_parquet_bytes( + table, write_size_statistics="columnchunk", write_page_index=True)) + pcc_idx = len(_write_parquet_bytes( + table, write_size_statistics="pageandcolumnchunk", + write_page_index=True)) + assert cc_idx < pcc_idx + + # Without the page index the two levels are indistinguishable. + pcc_noidx = len(_write_parquet_bytes( + table, write_size_statistics="pageandcolumnchunk")) + assert cc == pcc_noidx + + +def test_write_size_statistics_default_writes_statistics(): + """The default (None) keeps Arrow C++'s default of writing full size + statistics: it is byte-identical to explicit "pageandcolumnchunk" and + larger than "none".""" + table = _size_statistics_table() + default = _write_parquet_bytes(table, write_page_index=True) + explicit = _write_parquet_bytes( + table, write_page_index=True, + write_size_statistics="pageandcolumnchunk") + off = _write_parquet_bytes( + table, write_page_index=True, write_size_statistics="none") + assert default == explicit + assert len(off) < len(default) + + def test_page_checksum_verification_write_table(tempdir): """Check that checksum verification works for datasets created with pq.write_table()"""