diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd index 9cdba15638e6..c08487aa6a31 100644 --- a/python/pyarrow/_parquet.pxd +++ b/python/pyarrow/_parquet.pxd @@ -58,7 +58,8 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( store_decimal_as_integer=*, use_content_defined_chunking=*, bloom_filter_options=*, - min_compression_savings=* + min_compression_savings=*, + size_statistics_level=* ) except * diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index 40c504ef3b04..a66b02a60855 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -2077,7 +2077,8 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( store_decimal_as_integer=False, use_content_defined_chunking=False, bloom_filter_options=None, - min_compression_savings=0.01) except *: + min_compression_savings=0.01, + size_statistics_level=None) except *: """General writer properties""" cdef: @@ -2287,6 +2288,19 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( else: props.disable_write_page_index() + # size statistics + + if size_statistics_level is not None: + if size_statistics_level == "none": + props.set_size_statistics_level(SizeStatisticsLevel_None) + elif size_statistics_level == "column": + props.set_size_statistics_level(SizeStatisticsLevel_ColumnChunk) + elif size_statistics_level == "page_and_column": + props.set_size_statistics_level(SizeStatisticsLevel_PageAndColumnChunk) + else: + raise ValueError( + f"Unsupported size statistics level: {size_statistics_level}") + properties = props.build() return properties @@ -2415,7 +2429,8 @@ cdef class ParquetWriter(_Weakrefable): use_content_defined_chunking=False, write_time_adjusted_to_utc=False, bloom_filter_options=None, - min_compression_savings=0.01): + min_compression_savings=0.01, + size_statistics_level=None): cdef: shared_ptr[WriterProperties] properties shared_ptr[ArrowWriterProperties] arrow_properties @@ -2454,6 +2469,7 @@ cdef class ParquetWriter(_Weakrefable): use_content_defined_chunking=use_content_defined_chunking, bloom_filter_options=bloom_filter_options, min_compression_savings=min_compression_savings, + size_statistics_level=size_statistics_level, ) 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 a7ad239feac6..79a57b2b6c39 100644 --- a/python/pyarrow/includes/libparquet.pxd +++ b/python/pyarrow/includes/libparquet.pxd @@ -523,6 +523,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: @@ -606,6 +607,12 @@ cdef extern from "parquet/arrow/schema.h" namespace "parquet::arrow" nogil: cdef extern from "parquet/properties.h" namespace "parquet" nogil: + cdef enum SizeStatisticsLevel "parquet::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 ab5eafe51718..fba857406756 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -1001,6 +1001,9 @@ def _sanitize_table(table, new_schema, flavor): levels. The default is 0.01, so compression must save at least 1%. Values above this default require ``data_page_version="2.0"``. Using a larger value with V1 or a number outside the allowed range raises ``ValueError``. +size_statistics_level : {"none", "column", "page_and_column"}, default None + Set the level of size statistics written to the file metadata. If None, + the default size statistics behavior is used. """ _parquet_writer_example_doc = """\ @@ -1065,7 +1068,7 @@ class ParquetWriter: schema : pyarrow.Schema {_parquet_writer_arg_docs} writer_engine_version : unused -**options : dict +**kwargs : dict If options contains a key `metadata_collector` then the corresponding value is assumed to be a list (or any object with `.append` method) that will be filled with the file metadata instance @@ -1083,6 +1086,9 @@ def __init__(self, where, schema, filesystem=None, compression='snappy', write_statistics=True, use_deprecated_int96_timestamps=None, + coerce_timestamps=None, + allow_truncated_timestamps=False, + data_page_size=None, compression_level=None, use_byte_stream_split=False, column_encoding=None, @@ -1102,7 +1108,8 @@ def __init__(self, where, schema, filesystem=None, bloom_filter_options=None, use_content_defined_chunking=False, min_compression_savings=0.01, - **options): + size_statistics_level=None, + **kwargs): if use_deprecated_int96_timestamps is None: # Use int96 timestamps for Spark if flavor is not None and 'spark' in flavor: @@ -1110,6 +1117,9 @@ def __init__(self, where, schema, filesystem=None, else: use_deprecated_int96_timestamps = False + size_statistics_level = kwargs.pop('size_statistics_level', + size_statistics_level) + self.flavor = flavor if flavor is not None: schema, self.schema_changed = _sanitize_schema(schema, flavor) @@ -1132,7 +1142,7 @@ def __init__(self, where, schema, filesystem=None, path, compression=None) else: sink = where - self._metadata_collector = options.pop('metadata_collector', None) + self._metadata_collector = kwargs.pop('metadata_collector', None) engine_version = 'V2' self.writer = _parquet.ParquetWriter( sink, schema, @@ -1141,6 +1151,9 @@ def __init__(self, where, schema, filesystem=None, use_dictionary=use_dictionary, write_statistics=write_statistics, use_deprecated_int96_timestamps=use_deprecated_int96_timestamps, + coerce_timestamps=coerce_timestamps, + data_page_size=data_page_size, + allow_truncated_timestamps=allow_truncated_timestamps, compression_level=compression_level, use_byte_stream_split=use_byte_stream_split, column_encoding=column_encoding, @@ -1160,7 +1173,8 @@ def __init__(self, where, schema, filesystem=None, bloom_filter_options=bloom_filter_options, use_content_defined_chunking=use_content_defined_chunking, min_compression_savings=min_compression_savings, - **options) + size_statistics_level=size_statistics_level, + **kwargs) self.is_open = True def __del__(self): diff --git a/python/pyarrow/tests/parquet/test_datetime.py b/python/pyarrow/tests/parquet/test_datetime.py index b89fd97cb91e..42d76932ba39 100644 --- a/python/pyarrow/tests/parquet/test_datetime.py +++ b/python/pyarrow/tests/parquet/test_datetime.py @@ -158,6 +158,47 @@ def test_coerce_timestamps_truncated(tempdir): tm.assert_frame_equal(df_expected, df_ms) +def test_parquet_writer_accepts_timestamp_options(tempdir): + table = pa.table({ + 'ts': [datetime.datetime(2020, 1, 1), datetime.datetime(2020, 1, 2)] + }) + filename = tempdir / 'writer_kwargs.parquet' + + writer = pq.ParquetWriter( + filename, + table.schema, + coerce_timestamps='ms', + allow_truncated_timestamps=True, + size_statistics_level='page_and_column', + ) + writer.write_table(table) + writer.close() + + result = pq.read_table(filename) + assert result.schema.field('ts').type == pa.timestamp('ms') + + +def test_parquet_writer_accepts_size_statistics_level(tempdir): + table = pa.table({ + 'a': [1, 2], + 'ts': [datetime.datetime(2020, 1, 1), datetime.datetime(2020, 1, 2)] + }) + + writer_filename = tempdir / 'writer_size_statistics.parquet' + writer = pq.ParquetWriter( + writer_filename, + table.schema, + size_statistics_level='page_and_column', + ) + writer.write_table(table) + writer.close() + assert pq.read_table(writer_filename).num_rows == 2 + + write_table_filename = tempdir / 'write_table_size_statistics.parquet' + pq.write_table(table, write_table_filename, size_statistics_level='column') + assert pq.read_table(write_table_filename).num_rows == 2 + + @pytest.mark.pandas def test_date_time_types(tempdir): t1 = pa.date32()