diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index 3296af62f0c4..bba5efcdba3d 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -1057,8 +1057,12 @@ void ColumnWriterImpl::BuildDataPageV2(int64_t definition_levels_rle_size, bool page_is_compressed = false; if (pager_->has_compressor() && values->size() > 0) { pager_->Compress(*values, compressor_temp_buffer_.get()); - if (compressor_temp_buffer_->size() < values->size()) { - page_is_compressed = true; + if (const auto minimum = properties_->min_space_savings()) { + const double savings = + 1.0 - static_cast(compressor_temp_buffer_->size()) / values->size(); + page_is_compressed = savings >= *minimum; + } else { + page_is_compressed = compressor_temp_buffer_->size() < values->size(); } } std::shared_ptr compressed_values = diff --git a/cpp/src/parquet/column_writer_test.cc b/cpp/src/parquet/column_writer_test.cc index 69ec07c1b78f..02d45e197b37 100644 --- a/cpp/src/parquet/column_writer_test.cc +++ b/cpp/src/parquet/column_writer_test.cc @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. +#include #include +#include #include #include @@ -112,7 +114,8 @@ class TestPrimitiveWriter : public PrimitiveTypedTest { const ParquetVersion::type version = ParquetVersion::PARQUET_1_0, const ParquetDataPageVersion data_page_version = ParquetDataPageVersion::V1, bool enable_checksum = false, int64_t page_size = kDefaultDataPageSize, - int64_t max_rows_per_page = kDefaultMaxRowsPerPage) { + int64_t max_rows_per_page = kDefaultMaxRowsPerPage, + std::optional min_space_savings = {}) { sink_ = CreateOutputStream(); WriterProperties::Builder wp_builder; wp_builder.version(version)->data_page_version(data_page_version); @@ -130,6 +133,7 @@ class TestPrimitiveWriter : public PrimitiveTypedTest { wp_builder.max_statistics_size(column_properties.max_statistics_size()); wp_builder.data_pagesize(page_size); wp_builder.max_rows_per_page(max_rows_per_page); + wp_builder.min_space_savings(min_space_savings); writer_properties_ = wp_builder.build(); metadata_ = ColumnChunkMetaDataBuilder::Make(writer_properties_, this->descr_); @@ -2084,6 +2088,99 @@ TEST_F(TestValuesWriterInt32Type, AvoidCompressedInDataPageV2) { verify_only_one_uncompressed_page(/*total_num_values=*/1); } } + +TEST_F(TestValuesWriterInt32Type, MinSpaceSavingsDataPageV2) { + this->SetUpSchema(Repetition::OPTIONAL); + this->descr_ = this->schema_.Column(0); + this->GenerateData(SMALL_SIZE); + std::fill(this->values_.begin(), this->values_.end(), 0); + + // Calculate the actual savings to test the exact threshold. + ASSERT_OK_AND_ASSIGN(auto codec, Codec::Create(Compression::ZSTD)); + const int64_t original_size = SMALL_SIZE * sizeof(int32_t); + std::vector input(original_size, 0); + std::vector compressed(codec->MaxCompressedLen(original_size, input.data())); + ASSERT_OK_AND_ASSIGN(auto compressed_size, + codec->Compress(original_size, input.data(), compressed.size(), + compressed.data())); + + const double savings = 1.0 - static_cast(compressed_size) / original_size; + ASSERT_GT(savings, 0.0); + ASSERT_LT(savings, 1.0); + + const std::vector, bool>> cases = { + {std::nullopt, true}, + {0.0, true}, + {std::nextafter(savings, 0.0), true}, + {savings, true}, + {std::nextafter(savings, 1.0), false}, + {1.0, false}, + }; + + ColumnProperties column_properties; + column_properties.set_compression(Compression::ZSTD); + + for (const auto& [minimum, expected_compressed] : cases) { + auto writer = this->BuildWriter( + SMALL_SIZE, column_properties, ParquetVersion::PARQUET_2_LATEST, + ParquetDataPageVersion::V2, false, kDefaultDataPageSize, + kDefaultMaxRowsPerPage, minimum); + writer->WriteBatch(SMALL_SIZE, this->def_levels_.data(), nullptr, + this->values_ptr_); + writer->Close(); + ASSERT_OK_AND_ASSIGN(auto buffer, this->sink_->Finish()); + + auto page_reader = PageReader::Open( + std::make_shared<::arrow::io::BufferReader>(buffer), SMALL_SIZE, + Compression::ZSTD, default_reader_properties(), *this->descr_); + auto page = page_reader->NextPage(); + ASSERT_NE(page, nullptr); + ASSERT_EQ(PageType::DATA_PAGE_V2, page->type()); + auto data_page = std::static_pointer_cast(page); + ASSERT_EQ(expected_compressed, data_page->is_compressed()); + ASSERT_EQ(page_reader->NextPage(), nullptr); + + // Verify that the data still reads correctly. + auto read_pages = PageReader::Open( + std::make_shared<::arrow::io::BufferReader>(buffer), SMALL_SIZE, + Compression::ZSTD, default_reader_properties(), *this->descr_); + auto reader = std::static_pointer_cast( + ColumnReader::Make(this->descr_, std::move(read_pages))); + std::vector values(SMALL_SIZE); + std::vector definitions(SMALL_SIZE); + int64_t values_read = 0; + ASSERT_EQ(SMALL_SIZE, + reader->ReadBatch(SMALL_SIZE, definitions.data(), nullptr, + values.data(), &values_read)); + ASSERT_EQ(SMALL_SIZE, values_read); + ASSERT_EQ(this->values_, values); + ASSERT_EQ(this->def_levels_, definitions); + } +} + +TEST_F(TestValuesWriterInt32Type, MinSpaceSavingsIgnoredForDataPageV1) { + this->GenerateData(SMALL_SIZE); + ColumnProperties column_properties; + column_properties.set_compression(Compression::ZSTD); + std::shared_ptr default_output; + + for (std::optional minimum : + {std::optional{}, std::optional{1.0}}) { + auto writer = this->BuildWriter( + SMALL_SIZE, column_properties, ParquetVersion::PARQUET_2_LATEST, + ParquetDataPageVersion::V1, false, kDefaultDataPageSize, + kDefaultMaxRowsPerPage, minimum); + writer->WriteBatch(SMALL_SIZE, nullptr, nullptr, this->values_ptr_); + writer->Close(); + ASSERT_OK_AND_ASSIGN(auto buffer, this->sink_->Finish()); + + if (!minimum) { + default_output = buffer; + } else { + ASSERT_TRUE(default_output->Equals(*buffer)); + } + } +} #endif // Test writing and reading geometry columns diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index 89f2be42bffd..07e4f6be0cb1 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -403,6 +403,7 @@ class PARQUET_EXPORT WriterProperties { max_rows_per_page_(properties.max_rows_per_page()), version_(properties.version()), data_page_version_(properties.data_page_version()), + min_space_savings_(properties.min_space_savings()), created_by_(properties.created_by()), store_decimal_as_integer_(properties.store_decimal_as_integer()), page_checksum_enabled_(properties.page_checksum_enabled()), @@ -529,6 +530,18 @@ class PARQUET_EXPORT WriterProperties { return this; } + /// Minimum space savings required to compress V2 data page values. + /// + /// Savings are (1.0 - compressed_size / uncompressed_size), excluding + /// repetition and definition levels. For example, 0.1 requires 10% savings. + /// If unset, keep compression whenever it makes the values smaller. + /// Ignored for V1 pages and when compression is disabled. + /// Values outside [0, 1], including NaN, cause build() to throw. + Builder* min_space_savings(std::optional min_space_savings) { + min_space_savings_ = min_space_savings; + return this; + } + /// Specify the Parquet file version. /// Default PARQUET_2_6. Builder* version(ParquetVersion::type version) { @@ -883,6 +896,10 @@ class PARQUET_EXPORT WriterProperties { /// \brief Build the WriterProperties with the builder parameters. /// \return The WriterProperties defined by the builder. std::shared_ptr build() { + if (min_space_savings_ && + !(*min_space_savings_ >= 0.0 && *min_space_savings_ <= 1.0)) { + throw ParquetException("min_space_savings must be in range [0, 1]"); + } std::unordered_map column_properties; auto get = [&](const std::string& key) -> ColumnProperties& { auto it = column_properties.find(key); @@ -918,7 +935,7 @@ class PARQUET_EXPORT WriterProperties { pagesize_, max_rows_per_page_, version_, created_by_, page_checksum_enabled_, size_statistics_level_, std::move(file_encryption_properties_), default_column_properties_, column_properties, data_page_version_, - store_decimal_as_integer_, std::move(sorting_columns_), + min_space_savings_, store_decimal_as_integer_, std::move(sorting_columns_), content_defined_chunking_enabled_, content_defined_chunking_options_)); } @@ -933,6 +950,7 @@ class PARQUET_EXPORT WriterProperties { int64_t max_rows_per_page_; ParquetVersion::type version_; ParquetDataPageVersion data_page_version_; + std::optional min_space_savings_ = {}; std::string created_by_; bool store_decimal_as_integer_; bool page_checksum_enabled_; @@ -973,6 +991,9 @@ class PARQUET_EXPORT WriterProperties { return parquet_data_page_version_; } + /// Minimum space savings required to compress V2 data page values, if set. + std::optional min_space_savings() const { return min_space_savings_; } + inline ParquetVersion::type version() const { return parquet_version_; } inline std::string created_by() const { return parquet_created_by_; } @@ -1100,7 +1121,8 @@ class PARQUET_EXPORT WriterProperties { std::shared_ptr file_encryption_properties, const ColumnProperties& default_column_properties, const std::unordered_map& column_properties, - ParquetDataPageVersion data_page_version, bool store_short_decimal_as_integer, + ParquetDataPageVersion data_page_version, std::optional min_space_savings, + bool store_short_decimal_as_integer, std::vector sorting_columns, bool content_defined_chunking_enabled, CdcOptions content_defined_chunking_options) : pool_(pool), @@ -1110,6 +1132,7 @@ class PARQUET_EXPORT WriterProperties { pagesize_(pagesize), max_rows_per_page_(max_rows_per_page), parquet_data_page_version_(data_page_version), + min_space_savings_(min_space_savings), parquet_version_(version), parquet_created_by_(created_by), store_decimal_as_integer_(store_short_decimal_as_integer), @@ -1129,6 +1152,7 @@ class PARQUET_EXPORT WriterProperties { int64_t pagesize_; int64_t max_rows_per_page_; ParquetDataPageVersion parquet_data_page_version_; + std::optional min_space_savings_; ParquetVersion::type parquet_version_; std::string parquet_created_by_; bool store_decimal_as_integer_; diff --git a/cpp/src/parquet/properties_test.cc b/cpp/src/parquet/properties_test.cc index 2f22844f8bbf..9b84c78046e6 100644 --- a/cpp/src/parquet/properties_test.cc +++ b/cpp/src/parquet/properties_test.cc @@ -80,6 +80,31 @@ TEST(TestWriterProperties, DefaultCompression) { ::arrow::util::kUseDefaultCompressionLevel); } +TEST(TestWriterProperties, MinSpaceSavings) { + WriterProperties::Builder builder; + ASSERT_FALSE(builder.build()->min_space_savings().has_value()); + + for (double savings : {0.0, 0.1, 1.0}) { + auto properties = builder.min_space_savings(savings)->build(); + ASSERT_EQ(std::optional(savings), properties->min_space_savings()); + + auto copied = WriterProperties::Builder(*properties).build(); + ASSERT_EQ(properties->min_space_savings(), copied->min_space_savings()); + } + + auto properties = builder.min_space_savings(std::nullopt)->build(); + ASSERT_FALSE(properties->min_space_savings().has_value()); +} + +TEST(TestWriterProperties, InvalidMinSpaceSavings) { + for (double savings : {-0.1, 1.1, std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}) { + EXPECT_THROW(WriterProperties::Builder().min_space_savings(savings)->build(), + ParquetException); + } +} + TEST(TestWriterProperties, AdvancedHandling) { WriterProperties::Builder builder; builder.compression("gzip", Compression::GZIP); diff --git a/python/pyarrow/_dataset_parquet.pyx b/python/pyarrow/_dataset_parquet.pyx index 43eca8957be8..ffff69f4bcfe 100644 --- a/python/pyarrow/_dataset_parquet.pyx +++ b/python/pyarrow/_dataset_parquet.pyx @@ -655,6 +655,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): ), column_encoding=self._properties["column_encoding"], data_page_version=self._properties["data_page_version"], + min_space_savings=self._properties["min_space_savings"], encryption_properties=self._properties["encryption_properties"], write_batch_size=self._properties["write_batch_size"], dictionary_pagesize_limit=self._properties["dictionary_pagesize_limit"], @@ -707,6 +708,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): use_byte_stream_split=False, column_encoding=None, data_page_version="1.0", + min_space_savings=None, use_deprecated_int96_timestamps=False, coerce_timestamps=None, allow_truncated_timestamps=False, diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd index 36fc2ccf2f33..28a5c676ec98 100644 --- a/python/pyarrow/_parquet.pxd +++ b/python/pyarrow/_parquet.pxd @@ -57,7 +57,8 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( sorting_columns=*, store_decimal_as_integer=*, use_content_defined_chunking=*, - bloom_filter_options=* + bloom_filter_options=*, + min_space_savings=* ) except * diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index 2621c3bd6d61..640a69a2453d 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -2076,13 +2076,21 @@ cdef shared_ptr[WriterProperties] _create_writer_properties( sorting_columns=None, store_decimal_as_integer=False, use_content_defined_chunking=False, - bloom_filter_options=None) except *: + bloom_filter_options=None, + min_space_savings=None) except *: """General writer properties""" cdef: shared_ptr[WriterProperties] properties WriterProperties.Builder props CdcOptions cdc_options + double savings + + if min_space_savings is not None: + savings = min_space_savings + if not (0.0 <= savings <= 1.0): + raise ValueError("min_space_savings must be in range [0, 1]") + props.min_space_savings(optional[double](savings)) # data_page_version @@ -2405,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, + min_space_savings=None): cdef: shared_ptr[WriterProperties] properties shared_ptr[ArrowWriterProperties] arrow_properties @@ -2442,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, + min_space_savings=min_space_savings, ) 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 3a5f91716c17..c95d81a26761 100644 --- a/python/pyarrow/includes/libparquet.pxd +++ b/python/pyarrow/includes/libparquet.pxd @@ -485,6 +485,7 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil: cdef cppclass WriterProperties: cppclass Builder: Builder* data_page_version(ParquetDataPageVersion version) + Builder* min_space_savings(optional[double] min_space_savings) Builder* version(ParquetVersion version) Builder* compression(ParquetCompression codec) Builder* compression(const c_string& path, diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index 4de88fe22997..bccf47b917ed 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -1090,6 +1090,7 @@ def __init__(self, where, schema, filesystem=None, max_rows_per_page=None, bloom_filter_options=None, use_content_defined_chunking=False, + min_space_savings=None, **options): if use_deprecated_int96_timestamps is None: # Use int96 timestamps for Spark @@ -1147,6 +1148,7 @@ def __init__(self, where, schema, filesystem=None, max_rows_per_page=max_rows_per_page, bloom_filter_options=bloom_filter_options, use_content_defined_chunking=use_content_defined_chunking, + min_space_savings=min_space_savings, **options) self.is_open = True @@ -2039,6 +2041,7 @@ def write_table(table, where, row_group_size=None, version='2.6', max_rows_per_page=None, bloom_filter_options=None, use_content_defined_chunking=False, + min_space_savings=None, **kwargs): # Implementor's note: when adding keywords here / updating defaults, also # update it in write_to_dataset and _dataset_parquet.pyx ParquetFileWriteOptions @@ -2074,6 +2077,7 @@ def write_table(table, where, row_group_size=None, version='2.6', max_rows_per_page=max_rows_per_page, bloom_filter_options=bloom_filter_options, use_content_defined_chunking=use_content_defined_chunking, + min_space_savings=min_space_savings, **kwargs) as writer: writer.write_table(table, row_group_size=row_group_size) except Exception: diff --git a/python/pyarrow/tests/parquet/test_parquet_writer.py b/python/pyarrow/tests/parquet/test_parquet_writer.py index a49441f09f45..250ff05a3d06 100644 --- a/python/pyarrow/tests/parquet/test_parquet_writer.py +++ b/python/pyarrow/tests/parquet/test_parquet_writer.py @@ -41,6 +41,50 @@ pytestmark = pytest.mark.parquet +@pytest.mark.parametrize("minimum", [None, 0.0, 0.1, 1.0]) +@pytest.mark.parametrize("page_version", ["1.0", "2.0"]) +@pytest.mark.parametrize("compression", ["none", "snappy"]) +@pytest.mark.parametrize("use_writer", [False, True]) +def test_min_space_savings( + minimum, page_version, compression, use_writer): + table = pa.table({ + "value": pa.array([0.0] * 1000, type=pa.float32()) + }) + sink = pa.BufferOutputStream() + options = dict( + compression=compression, + use_dictionary=False, + data_page_version=page_version, + min_space_savings=minimum, + ) + + if use_writer: + with pq.ParquetWriter(sink, table.schema, **options) as writer: + writer.write_table(table) + else: + pq.write_table(table, sink, **options) + + assert pq.read_table(sink.getvalue()).equals(table) + + +@pytest.mark.parametrize( + "minimum", + [-0.1, 1.1, float("inf"), -float("inf"), float("nan")] +) +@pytest.mark.parametrize("use_writer", [False, True]) +def test_invalid_min_space_savings(minimum, use_writer): + table = pa.table({"value": [1.0]}) + sink = pa.BufferOutputStream() + + with pytest.raises(ValueError, match="min_space_savings"): + if use_writer: + pq.ParquetWriter( + sink, table.schema, min_space_savings=minimum) + else: + pq.write_table( + table, sink, min_space_savings=minimum) + + @pytest.mark.pandas def test_parquet_incremental_file_build(tempdir): df = _test_dataframe(100) diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 0a94c0bd9875..7daf2dfb2558 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -5949,6 +5949,27 @@ def test_write_dataset_content_defined_chunking(tempdir, cdc): assert ds.dataset(tempdir, format=fmt).to_table().equals(expected_table) +@pytest.mark.parquet +@pytest.mark.parametrize("minimum", [None, 0.0, 0.1, 1.0]) +def test_parquet_min_space_savings(tempdir, minimum): + import pyarrow.parquet as pq + + table = pa.table({ + "value": pa.array([0.0] * 1000, type=pa.float32()) + }) + pq.write_to_dataset( + table, + tempdir, + compression="snappy", + use_dictionary=False, + data_page_version="2.0", + min_space_savings=minimum, + ) + + result = ds.dataset(tempdir, format="parquet").to_table() + assert result.equals(table) + + def test_make_write_options_error(): # GH-39440: calling make_write_options as a static class method msg_1 = ("make_write_options() should be called on an "