From fec0099864461c55762f932af87c6c515270bc6f Mon Sep 17 00:00:00 2001 From: bodduv Date: Sat, 3 Oct 2026 14:26:50 +0200 Subject: [PATCH] DX-124448: [Parquet] Add first-page dictionary benefit check The Parquet writer keeps dictionary encoding until the dictionary reaches dictionary_pagesize_limit, even when the dictionary saves nothing. On high-cardinality columns this dictionary-encodes about half of each column chunk (a 2 MiB dictionary of 16-byte strings holds about 105K values), costing CPU and about 10% more bytes than PLAIN. parquet-java decides on the first data page of each column chunk: it keeps the dictionary only if the page's encoded indices plus the dictionary are smaller than the page's values in PLAIN encoding. WriterProperties::Builder::enable_dictionary_benefit_check() applies that rule. When the first page fails it, the writer falls back to PLAIN through the existing FallbackToPlainEncoding(), so the dictionary page and the first page stay dictionary-encoded. The check is off by default. A column chunk that ends before its first page is cut, and values written as an Arrow DictionaryArray, are not checked. Co-Authored-By: Claude Opus 5.5 --- cpp/src/parquet/column_writer.cc | 65 ++++++++++++++++++-- cpp/src/parquet/column_writer_test.cc | 86 +++++++++++++++++++++++++++ cpp/src/parquet/properties.h | 41 +++++++++++-- cpp/src/parquet/properties_test.cc | 13 ++++ 4 files changed, 196 insertions(+), 9 deletions(-) diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index 94b67dfa807e..bf8e90f5b851 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -896,6 +896,9 @@ class ColumnWriterImpl { // Flag to infer if dictionary encoding has fallen back to PLAIN bool fallback_; + // Size of the encoded values (without levels) of the most recent data page + int64_t last_page_values_size_ = 0; + ::arrow::BufferBuilder definition_levels_sink_; ::arrow::BufferBuilder repetition_levels_sink_; @@ -964,6 +967,7 @@ void ColumnWriterImpl::AddDataPage() { int64_t repetition_levels_rle_size = 0; std::shared_ptr values = GetValuesBuffer(); + last_page_values_size_ = values->size(); bool is_v1_data_page = properties_->data_page_version() == ParquetDataPageVersion::V1; if (descr_->max_definition_level() > 0) { @@ -1641,6 +1645,11 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, // which case we call back to the dense write path) std::shared_ptr<::arrow::Array> preserved_dictionary_; + // State of the first-page dictionary benefit check, see + // WriterProperties::dictionary_benefit_check_enabled(). + bool dictionary_benefit_checked_ = false; + int64_t first_page_plain_size_ = 0; + int64_t WriteLevels(int64_t num_levels, const int16_t* def_levels, const int16_t* rep_levels) { // Update histograms now, to maximize cache efficiency. @@ -1798,11 +1807,17 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, } // Update the unencoded data bytes for ByteArray only per the specification. - void UpdateUnencodedDataBytes() const { + void UpdateUnencodedDataBytes() { if constexpr (std::is_same_v) { - if (page_size_statistics_ != nullptr) { - page_size_statistics_->IncrementUnencodedByteArrayDataBytes( - current_encoder_->ReportUnencodedDataBytes()); + const bool checking_dictionary_benefit = IsCheckingDictionaryBenefit(); + if (page_size_statistics_ != nullptr || checking_dictionary_benefit) { + const int64_t unencoded_bytes = current_encoder_->ReportUnencodedDataBytes(); + if (page_size_statistics_ != nullptr) { + page_size_statistics_->IncrementUnencodedByteArrayDataBytes(unencoded_bytes); + } + if (checking_dictionary_benefit) { + first_page_plain_size_ += unencoded_bytes; + } } } } @@ -1816,7 +1831,43 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, if (check_page_limit && (current_encoder_->EstimatedDataEncodedSize() >= properties_->data_pagesize() || num_buffered_rows_ >= properties_->max_rows_per_page())) { + const bool first_dictionary_page = IsCheckingDictionaryBenefit(); AddDataPage(); + if (first_dictionary_page) { + CheckDictionaryBenefit(); + } + } + } + + // Whether the values written so far belong to the first data page of a + // dictionary-encoded column chunk that the benefit check has not judged yet. + bool IsCheckingDictionaryBenefit() const { + return properties_->dictionary_benefit_check_enabled() && has_dictionary_ && + !fallback_ && !dictionary_benefit_checked_; + } + + // Keeps the dictionary only if the first page's encoded indices plus the + // dictionary are smaller than the page in PLAIN encoding, as parquet-java does. + void CheckDictionaryBenefit() { + dictionary_benefit_checked_ = true; + if (last_page_values_size_ + current_dict_encoder_->dict_encoded_size() >= + first_page_plain_size_) { + FallbackToPlainEncoding(); + } + } + + // Adds the PLAIN-encoded size of non-null values written to the first page. + // Byte-array lengths are added separately by UpdateUnencodedDataBytes(). + void AddFirstPagePlainSize(int64_t num_non_null) { + if (!IsCheckingDictionaryBenefit()) { + return; + } + if constexpr (std::is_same_v) { + first_page_plain_size_ += num_non_null * static_cast(sizeof(uint32_t)); + } else if constexpr (std::is_same_v) { + first_page_plain_size_ += num_non_null * descr_->type_length(); + } else { + first_page_plain_size_ += num_non_null * static_cast(sizeof(T)); } } @@ -1860,6 +1911,7 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, page_statistics_->Update(values, num_values, num_nulls); } + AddFirstPagePlainSize(num_values); UpdateUnencodedDataBytes(); if constexpr (std::is_same::value) { @@ -1896,6 +1948,7 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, num_spaced_values, num_values, num_nulls); } + AddFirstPagePlainSize(num_values); UpdateUnencodedDataBytes(); if constexpr (std::is_same::value) { @@ -1945,6 +1998,9 @@ Status TypedColumnWriterImpl::WriteArrowDictionary( return WriteDense(); } + // Values arrive as dictionary indices, so the first page's PLAIN size is unknown. + dictionary_benefit_checked_ = true; + auto dict_encoder = dynamic_cast*>(current_encoder_.get()); const auto& data = checked_cast(array); std::shared_ptr<::arrow::Array> dictionary = data.dictionary(); @@ -2471,6 +2527,7 @@ Status TypedColumnWriterImpl::WriteArrowDense( page_statistics_->IncrementNumValues(non_null); } + AddFirstPagePlainSize(non_null); UpdateUnencodedDataBytes(); if (chunk_geospatial_statistics_ != nullptr) { diff --git a/cpp/src/parquet/column_writer_test.cc b/cpp/src/parquet/column_writer_test.cc index dedf25abcabc..b7b69dd35813 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 @@ -125,6 +127,9 @@ class TestPrimitiveWriter : public PrimitiveTypedTest { if (enable_checksum) { wp_builder.enable_page_checksum(); } + if (enable_dictionary_benefit_check_) { + wp_builder.enable_dictionary_benefit_check(); + } 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); @@ -275,6 +280,46 @@ class TestPrimitiveWriter : public PrimitiveTypedTest { } } + // Writes VERY_LARGE_SIZE values in pages of 1,000 rows with the dictionary benefit + // check enabled and returns the number of data pages per encoding. + std::map WriteWithDictionaryBenefitCheck( + ParquetVersion::type version, ParquetDataPageVersion data_page_version, + int64_t distinct_values) { + this->GenerateData(VERY_LARGE_SIZE); + if (distinct_values < VERY_LARGE_SIZE) { + for (int64_t i = distinct_values; i < VERY_LARGE_SIZE; i++) { + this->values_[i] = this->values_[i % distinct_values]; + } + } + ColumnProperties column_properties; + column_properties.set_dictionary_enabled(true); + column_properties.set_encoding(version == ParquetVersion::PARQUET_1_0 + ? Encoding::PLAIN_DICTIONARY + : Encoding::RLE_DICTIONARY); + enable_dictionary_benefit_check_ = true; + auto writer = this->BuildWriter(VERY_LARGE_SIZE, column_properties, version, + data_page_version, /*enable_checksum=*/false, + kDefaultDataPageSize, /*max_rows_per_page=*/1000); + enable_dictionary_benefit_check_ = false; + writer->WriteBatch(this->values_.size(), nullptr, nullptr, this->values_ptr_); + writer->Close(); + + this->SetupValuesOut(VERY_LARGE_SIZE); + this->ReadColumnFully(); + EXPECT_EQ(VERY_LARGE_SIZE, this->values_read_); + this->values_.resize(VERY_LARGE_SIZE); + EXPECT_EQ(this->values_, this->values_out_); + + std::map data_pages; + for (const auto& stats : this->metadata_encoding_stats()) { + if (stats.page_type == PageType::DATA_PAGE || + stats.page_type == PageType::DATA_PAGE_V2) { + data_pages[stats.encoding] += stats.count; + } + } + return data_pages; + } + void WriteRequiredWithSettings(Encoding::type encoding, Compression::type compression, bool enable_dictionary, bool enable_statistics, int compression_level, int64_t num_rows, @@ -438,6 +483,7 @@ class TestPrimitiveWriter : public PrimitiveTypedTest { std::shared_ptr<::arrow::io::BufferOutputStream> sink_; std::shared_ptr writer_properties_; std::vector> data_buffer_; + bool enable_dictionary_benefit_check_ = false; }; template @@ -799,6 +845,46 @@ TYPED_TEST(TestPrimitiveWriter, DictionaryFallbackVersion1_0) { ParquetDataPageVersion::V1); } +// The first page of unique values gains nothing from a dictionary, so the writer falls +// back to PLAIN after that page instead of at the dictionary size limit. +TYPED_TEST(TestPrimitiveWriter, DictionaryBenefitCheckFallsBackAfterFirstPage) { + if (this->type_num() == Type::BOOLEAN) { + GTEST_SKIP() << "BOOLEAN is never dictionary-encoded"; + } + for (auto [version, data_page_version, dictionary_encoding] : + {std::tuple{ParquetVersion::PARQUET_1_0, ParquetDataPageVersion::V1, + Encoding::PLAIN_DICTIONARY}, + std::tuple{ParquetVersion::PARQUET_2_6, ParquetDataPageVersion::V1, + Encoding::RLE_DICTIONARY}, + std::tuple{ParquetVersion::PARQUET_2_6, ParquetDataPageVersion::V2, + Encoding::RLE_DICTIONARY}}) { + auto data_pages = this->WriteWithDictionaryBenefitCheck( + version, data_page_version, /*distinct_values=*/VERY_LARGE_SIZE); + std::map expected{ + {dictionary_encoding, 1}, {Encoding::PLAIN, VERY_LARGE_SIZE / 1000 - 1}}; + ASSERT_EQ(expected, data_pages); + } +} + +// A low-cardinality first page shrinks under dictionary encoding, so the dictionary is +// kept for the whole column chunk. +TYPED_TEST(TestPrimitiveWriter, DictionaryBenefitCheckKeepsBeneficialDictionary) { + if (this->type_num() == Type::BOOLEAN) { + GTEST_SKIP() << "BOOLEAN is never dictionary-encoded"; + } + for (auto [version, data_page_version, dictionary_encoding] : + {std::tuple{ParquetVersion::PARQUET_1_0, ParquetDataPageVersion::V1, + Encoding::PLAIN_DICTIONARY}, + std::tuple{ParquetVersion::PARQUET_2_6, ParquetDataPageVersion::V2, + Encoding::RLE_DICTIONARY}}) { + auto data_pages = this->WriteWithDictionaryBenefitCheck(version, data_page_version, + /*distinct_values=*/10); + std::map expected{ + {dictionary_encoding, VERY_LARGE_SIZE / 1000}}; + ASSERT_EQ(expected, data_pages); + } +} + TYPED_TEST(TestPrimitiveWriter, DictionaryFallbackVersion2_0) { this->TestDictionaryFallbackEncoding(ParquetVersion::PARQUET_2_4, ParquetDataPageVersion::V1); diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index eb5aee29695e..fafb4fc72bb1 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -302,7 +302,8 @@ class PARQUET_EXPORT WriterProperties { page_checksum_enabled_(false), size_statistics_level_(DEFAULT_SIZE_STATISTICS_LEVEL), content_defined_chunking_enabled_(false), - content_defined_chunking_options_({}) {} + content_defined_chunking_options_({}), + dictionary_benefit_check_enabled_(false) {} explicit Builder(const WriterProperties& properties) : pool_(properties.memory_pool()), @@ -322,7 +323,9 @@ class PARQUET_EXPORT WriterProperties { content_defined_chunking_enabled_( properties.content_defined_chunking_enabled()), content_defined_chunking_options_( - properties.content_defined_chunking_options()) { + properties.content_defined_chunking_options()), + dictionary_benefit_check_enabled_( + properties.dictionary_benefit_check_enabled()) { CopyColumnSpecificProperties(properties); } @@ -404,6 +407,24 @@ class PARQUET_EXPORT WriterProperties { return this; } + /// \brief Fall back to PLAIN when the first data page of a column chunk shows + /// that dictionary encoding does not reduce its size. + /// + /// The check keeps the dictionary only if the first page's encoded indices plus + /// the dictionary are smaller than the page's values would be in PLAIN encoding, + /// the rule parquet-java applies. Without it, the dictionary is abandoned only + /// when it reaches dictionary_pagesize_limit. Disabled by default. + Builder* enable_dictionary_benefit_check() { + dictionary_benefit_check_enabled_ = true; + return this; + } + + /// Disable the first-page dictionary benefit check (the default). + Builder* disable_dictionary_benefit_check() { + dictionary_benefit_check_enabled_ = false; + return this; + } + /// Specify the write batch size while writing batches of Arrow values /// into Parquet. Default 1024. Builder* write_batch_size(int64_t write_batch_size) { @@ -783,7 +804,8 @@ class PARQUET_EXPORT WriterProperties { size_statistics_level_, std::move(file_encryption_properties_), default_column_properties_, column_properties, data_page_version_, store_decimal_as_integer_, std::move(sorting_columns_), - content_defined_chunking_enabled_, content_defined_chunking_options_)); + content_defined_chunking_enabled_, content_defined_chunking_options_, + dictionary_benefit_check_enabled_)); } private: @@ -818,12 +840,19 @@ class PARQUET_EXPORT WriterProperties { bool content_defined_chunking_enabled_; CdcOptions content_defined_chunking_options_; + bool dictionary_benefit_check_enabled_; }; inline MemoryPool* memory_pool() const { return pool_; } inline int64_t dictionary_pagesize_limit() const { return dictionary_pagesize_limit_; } + /// \brief Whether the first data page of each column chunk decides if dictionary + /// encoding is kept. See Builder::enable_dictionary_benefit_check(). + inline bool dictionary_benefit_check_enabled() const { + return dictionary_benefit_check_enabled_; + } + inline int64_t write_batch_size() const { return write_batch_size_; } inline int64_t max_row_group_length() const { return max_row_group_length_; } @@ -954,7 +983,7 @@ class PARQUET_EXPORT WriterProperties { const std::unordered_map& column_properties, ParquetDataPageVersion data_page_version, bool store_short_decimal_as_integer, std::vector sorting_columns, bool content_defined_chunking_enabled, - CdcOptions content_defined_chunking_options) + CdcOptions content_defined_chunking_options, bool dictionary_benefit_check_enabled) : pool_(pool), dictionary_pagesize_limit_(dictionary_pagesize_limit), write_batch_size_(write_batch_size), @@ -972,7 +1001,8 @@ class PARQUET_EXPORT WriterProperties { default_column_properties_(default_column_properties), column_properties_(column_properties), content_defined_chunking_enabled_(content_defined_chunking_enabled), - content_defined_chunking_options_(content_defined_chunking_options) {} + content_defined_chunking_options_(content_defined_chunking_options), + dictionary_benefit_check_enabled_(dictionary_benefit_check_enabled) {} MemoryPool* pool_; int64_t dictionary_pagesize_limit_; @@ -996,6 +1026,7 @@ class PARQUET_EXPORT WriterProperties { bool content_defined_chunking_enabled_; CdcOptions content_defined_chunking_options_; + bool dictionary_benefit_check_enabled_; }; PARQUET_EXPORT const std::shared_ptr& default_writer_properties(); diff --git a/cpp/src/parquet/properties_test.cc b/cpp/src/parquet/properties_test.cc index 7be352aa5f11..e99c61df1238 100644 --- a/cpp/src/parquet/properties_test.cc +++ b/cpp/src/parquet/properties_test.cc @@ -48,6 +48,19 @@ TEST(TestWriterProperties, Basics) { ASSERT_EQ(ParquetVersion::PARQUET_2_6, props->version()); ASSERT_EQ(ParquetDataPageVersion::V1, props->data_page_version()); ASSERT_FALSE(props->page_checksum_enabled()); + ASSERT_FALSE(props->dictionary_benefit_check_enabled()); +} + +TEST(TestWriterProperties, DictionaryBenefitCheck) { + WriterProperties::Builder builder; + std::shared_ptr enabled = + builder.enable_dictionary_benefit_check()->build(); + ASSERT_TRUE(enabled->dictionary_benefit_check_enabled()); + ASSERT_TRUE( + WriterProperties::Builder(*enabled).build()->dictionary_benefit_check_enabled()); + ASSERT_FALSE(builder.disable_dictionary_benefit_check() + ->build() + ->dictionary_benefit_check_enabled()); } TEST(TestWriterProperties, DefaultCompression) {