PARQUET-3479: Add configuration to disable early dictionary compression check

parquet-column/src/main/java/org/apache/parquet/column/ParquetProperties.java

@@ -67,6 +67,7 @@ public class ParquetProperties {
   public static final boolean DEFAULT_STATISTICS_ENABLED = true;
   public static final boolean DEFAULT_SIZE_STATISTICS_ENABLED = true;
 
+  public static final long DEFAULT_DICTIONARY_CHECK_THRESHOLD_RAW_SIZE_BYTES = 0;
   public static final boolean DEFAULT_PAGE_WRITE_CHECKSUM_ENABLED = true;
 
   /**
@@ -131,6 +132,7 @@ public static WriterVersion fromString(String name) {
   private final int rowGroupRowCountLimit;
   private final int pageRowCountLimit;
   private final boolean pageWriteChecksumEnabled;
+  private final long dictionaryCheckThresholdRawSizeBytes;
   private final ColumnProperty<ByteStreamSplitMode> byteStreamSplitEnabled;
   private final Map<String, String> extraMetaData;
   private final ColumnProperty<Boolean> statistics;
@@ -163,6 +165,7 @@ private ParquetProperties(Builder builder) {
     this.rowGroupRowCountLimit = builder.rowGroupRowCountLimit;
     this.pageRowCountLimit = builder.pageRowCountLimit;
     this.pageWriteChecksumEnabled = builder.pageWriteChecksumEnabled;
+    this.dictionaryCheckThresholdRawSizeBytes = builder.dictionaryCheckThresholdRawSizeBytes;
     this.byteStreamSplitEnabled = builder.byteStreamSplitEnabled.build();
     this.extraMetaData = builder.extraMetaData;
     this.statistics = builder.statistics.build();
@@ -322,6 +325,17 @@ public boolean getPageWriteChecksumEnabled() {
     return pageWriteChecksumEnabled;
   }
 
+  /**
+   * Returns the byte threshold after which the dictionary compression check is performed.
+   * A value of 0 means check on the first page. Higher values delay the check until that
+   * many raw bytes have been accumulated across pages.
+   *
+   * @return the byte threshold for the dictionary compression check
+   */
+  public long getDictionaryCheckThresholdRawSizeBytes() {
+    return dictionaryCheckThresholdRawSizeBytes;
+  }
+
   public OptionalLong getBloomFilterNDV(ColumnDescriptor column) {
     Long ndv = bloomFilterNDVs.getValue(column);
     return ndv == null ? OptionalLong.empty() : OptionalLong.of(ndv);
@@ -415,6 +429,7 @@ public static class Builder {
     private int rowGroupRowCountLimit = DEFAULT_ROW_GROUP_ROW_COUNT_LIMIT;
     private int pageRowCountLimit = DEFAULT_PAGE_ROW_COUNT_LIMIT;
     private boolean pageWriteChecksumEnabled = DEFAULT_PAGE_WRITE_CHECKSUM_ENABLED;
+    private long dictionaryCheckThresholdRawSizeBytes = DEFAULT_DICTIONARY_CHECK_THRESHOLD_RAW_SIZE_BYTES;
     private final ColumnProperty.Builder<ByteStreamSplitMode> byteStreamSplitEnabled;
     private Map<String, String> extraMetaData = new HashMap<>();
     private final ColumnProperty.Builder<Boolean> statistics;
@@ -450,6 +465,7 @@ private Builder(ParquetProperties toCopy) {
       this.allocator = toCopy.allocator;
       this.pageRowCountLimit = toCopy.pageRowCountLimit;
       this.pageWriteChecksumEnabled = toCopy.pageWriteChecksumEnabled;
+      this.dictionaryCheckThresholdRawSizeBytes = toCopy.dictionaryCheckThresholdRawSizeBytes;
       this.bloomFilterNDVs = ColumnProperty.builder(toCopy.bloomFilterNDVs);
       this.bloomFilterFPPs = ColumnProperty.builder(toCopy.bloomFilterFPPs);
       this.bloomFilterEnabled = ColumnProperty.builder(toCopy.bloomFilterEnabled);
@@ -709,6 +725,20 @@ public Builder withPageWriteChecksumEnabled(boolean val) {
       return this;
     }
 
+    /**
+     * Set the raw data byte threshold after which the dictionary compression check is performed.
+     * A value of 0 means check on the first page. Higher values delay the check until that
+     * many raw bytes have been accumulated across pages.
+     *
+     * @param val byte threshold (default: 0)
+     * @return this builder for method chaining
+     */
+    public Builder withDictionaryCheckThresholdRawSizeBytes(long val) {
+      Preconditions.checkArgument(val >= 0, "dictionaryCheckThresholdRawSizeBytes must be >= 0");
+      this.dictionaryCheckThresholdRawSizeBytes = val;
+      return this;
+    }
+
     public Builder withExtraMetaData(Map<String, String> extraMetaData) {
       this.extraMetaData = extraMetaData;
       return this;

parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultValuesWriterFactory.java

@@ -111,7 +111,9 @@ static ValuesWriter dictWriterWithFallBack(
       ValuesWriter writerToFallBackTo) {
     if (parquetProperties.isDictionaryEnabled(path)) {
       return FallbackValuesWriter.of(
-          dictionaryWriter(path, parquetProperties, dictPageEncoding, dataPageEncoding), writerToFallBackTo);
+          dictionaryWriter(path, parquetProperties, dictPageEncoding, dataPageEncoding),
+          writerToFallBackTo,
+          parquetProperties.getDictionaryCheckThresholdRawSizeBytes());
     } else {
       return writerToFallBackTo;
     }

parquet-column/src/main/java/org/apache/parquet/column/values/fallback/FallbackValuesWriter.java

@@ -30,7 +30,12 @@ public class FallbackValuesWriter<I extends ValuesWriter & RequiresFallback, F e
 
   public static <I extends ValuesWriter & RequiresFallback, F extends ValuesWriter> FallbackValuesWriter<I, F> of(
       I initialWriter, F fallBackWriter) {
-    return new FallbackValuesWriter<>(initialWriter, fallBackWriter);
+    return new FallbackValuesWriter<>(initialWriter, fallBackWriter, /*dictionaryCheckThresholdRawSizeBytes=*/ 0);
+  }
+
+  public static <I extends ValuesWriter & RequiresFallback, F extends ValuesWriter> FallbackValuesWriter<I, F> of(
+      I initialWriter, F fallBackWriter, long dictionaryCheckThresholdRawSizeBytes) {
+    return new FallbackValuesWriter<>(initialWriter, fallBackWriter, dictionaryCheckThresholdRawSizeBytes);
   }
 
   /**
@@ -43,6 +48,17 @@ public static <I extends ValuesWriter & RequiresFallback, F extends ValuesWriter
   public final F fallBackWriter;
 
   private boolean fellBackAlready = false;
+  private boolean compressionChecked = false;
+  private final long dictionaryCheckThresholdRawSizeBytes;
+  /** Accumulates raw bytes across pages (only reset in resetDictionary) so the
+   * threshold check works even when individual pages are smaller than the threshold.
+   * Overflow is not a concern: a long would require writing over 9.2 exabytes to a single
+   * column chunk, which is physically impossible. */
+  private long cumulativeRawBytes = 0;
+  /** Accumulates dictionary-encoded size across pages (only reset in resetDictionary) so the
+   * compression decision compares like-scoped cumulative quantities — the column-chunk
+   * dictionary cost is amortized over all pages it covers, not charged against a single page. */
+  private long cumulativeEncodedBytes = 0;
 
   /**
    * writer currently written to
@@ -57,16 +73,16 @@ public static <I extends ValuesWriter & RequiresFallback, F extends ValuesWriter
    */
   private long rawDataByteSize = 0;
 
-  /**
-   * indicates if this is the first page being processed
-   */
-  private boolean firstPage = true;
-
   public FallbackValuesWriter(I initialWriter, F fallBackWriter) {
+    this(initialWriter, fallBackWriter, /*dictionaryCheckThresholdRawSizeBytes=*/ 0);
+  }
+
+  public FallbackValuesWriter(I initialWriter, F fallBackWriter, long dictionaryCheckThresholdRawSizeBytes) {
     super();
     this.initialWriter = initialWriter;
     this.fallBackWriter = fallBackWriter;
     this.currentWriter = initialWriter;
+    this.dictionaryCheckThresholdRawSizeBytes = dictionaryCheckThresholdRawSizeBytes;
   }
 
   @Override
@@ -79,16 +95,38 @@ public long getBufferedSize() {
 
   @Override
   public BytesInput getBytes() {
-    if (!fellBackAlready && firstPage) {
-      // we use the first page to decide if we're going to use this encoding
+    try {
+      cumulativeRawBytes = Math.addExact(cumulativeRawBytes, rawDataByteSize);
+    } catch (ArithmeticException e) {
+      // overflow, keep the previous value
+    }
+    if (!fellBackAlready && !compressionChecked && cumulativeRawBytes >= dictionaryCheckThresholdRawSizeBytes) {
+      compressionChecked = true;
       BytesInput bytes = initialWriter.getBytes();
-      if (!initialWriter.isCompressionSatisfying(rawDataByteSize, bytes.size())) {
+      try {
+        cumulativeEncodedBytes = Math.addExact(cumulativeEncodedBytes, bytes.size());
+      } catch (ArithmeticException e) {
+        // overflow, keep the previous value
+      }
+      // Compare cumulative raw vs cumulative encoded so the column-chunk dictionary
+      // (which is itself cumulative) is amortized over all pages it covers, not charged
+      // against a single page.
+      if (!initialWriter.isCompressionSatisfying(cumulativeRawBytes, cumulativeEncodedBytes)) {
         fallBack();
       } else {
         return bytes;
       }
     }
-    return currentWriter.getBytes();
+    BytesInput result = currentWriter.getBytes();
+    if (!fellBackAlready && !compressionChecked) {
+      // Accumulate dictionary-encoded size for pages flushed before the check fires.
+      try {
+        cumulativeEncodedBytes = Math.addExact(cumulativeEncodedBytes, result.size());
+      } catch (ArithmeticException e) {
+        // overflow, keep the previous value
+      }
+    }
+    return result;
   }
 
   @Override
@@ -103,7 +141,6 @@ public Encoding getEncoding() {
   @Override
   public void reset() {
     rawDataByteSize = 0;
-    firstPage = false;
     currentWriter.reset();
   }
 
@@ -131,8 +168,10 @@ public void resetDictionary() {
     }
     currentWriter = initialWriter;
     fellBackAlready = false;
+    compressionChecked = false;
+    cumulativeRawBytes = 0;
+    cumulativeEncodedBytes = 0;
     initialUsedAndHadDictionary = false;
-    firstPage = true;
   }
 
   @Override