DX-124448: [Parquet] Add first-page dictionary benefit check

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<Buffer> 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<T, ByteArray>) {
-      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<T, ByteArray>) {
+      first_page_plain_size_ += num_non_null * static_cast<int64_t>(sizeof(uint32_t));
+    } else if constexpr (std::is_same_v<T, FLBA>) {
+      first_page_plain_size_ += num_non_null * descr_->type_length();
+    } else {
+      first_page_plain_size_ += num_non_null * static_cast<int64_t>(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<T, ByteArray>::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<T, ByteArray>::value) {
@@ -1945,6 +1998,9 @@ Status TypedColumnWriterImpl<ParquetType>::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<DictEncoder<ParquetType>*>(current_encoder_.get());
   const auto& data = checked_cast<const ::arrow::DictionaryArray&>(array);
   std::shared_ptr<::arrow::Array> dictionary = data.dictionary();
@@ -2471,6 +2527,7 @@ Status TypedColumnWriterImpl<ByteArrayType>::WriteArrowDense(
       page_statistics_->IncrementNumValues(non_null);
     }
 
+    AddFirstPagePlainSize(non_null);
     UpdateUnencodedDataBytes();
 
     if (chunk_geospatial_statistics_ != nullptr) {

cpp/src/parquet/column_writer_test.cc

@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include <map>
 #include <memory>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -125,6 +127,9 @@ class TestPrimitiveWriter : public PrimitiveTypedTest<TestType> {
     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<TestType> {
     }
   }
 
+  // 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<Encoding::type, int32_t> 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<Encoding::type, int32_t> 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<TestType> {
   std::shared_ptr<::arrow::io::BufferOutputStream> sink_;
   std::shared_ptr<WriterProperties> writer_properties_;
   std::vector<std::vector<uint8_t>> data_buffer_;
+  bool enable_dictionary_benefit_check_ = false;
 };
 
 template <typename TestType>
@@ -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<Encoding::type, int32_t> 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<Encoding::type, int32_t> 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);

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<std::string, ColumnProperties>& column_properties,
       ParquetDataPageVersion data_page_version, bool store_short_decimal_as_integer,
       std::vector<SortingColumn> 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<WriterProperties>& default_writer_properties();

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<WriterProperties> 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) {