Parquet: Add opt-in uncompressed row group size tracking

core/src/main/java/org/apache/iceberg/TableProperties.java

@@ -176,6 +176,10 @@ private TableProperties() {}
       "write.delete.parquet.row-group-check-max-record-count";
   public static final int PARQUET_ROW_GROUP_CHECK_MAX_RECORD_COUNT_DEFAULT = 10000;
 
+  public static final String PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED =
+      "write.parquet.row-group-size-track-uncompressed";
+  public static final boolean PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED_DEFAULT = false;
+
   public static final String PARQUET_BLOOM_FILTER_MAX_BYTES =
       "write.parquet.bloom-filter-max-bytes";
   public static final int PARQUET_BLOOM_FILTER_MAX_BYTES_DEFAULT = 1024 * 1024;

docs/docs/configuration.md

@@ -44,6 +44,7 @@ Iceberg tables support table properties to configure table behavior, like the de
 | write.format.default                                | parquet                     | Default file format for the table; parquet, avro, or orc                                                                                                                                                                                           |
 | write.delete.format.default                         | data file format            | Default delete file format for the table; parquet, avro, or orc                                                                                                                                                                                    |
 | write.parquet.row-group-size-bytes                  | 134217728 (128 MB)          | Parquet row group size                                                                                                                                                                                                                             |
+| write.parquet.row-group-size-track-uncompressed     | false                       | Track raw data size to enforce the row group size target accurately with compressing codecs                                                                                                                                                         |
 | write.parquet.page-size-bytes                       | 1048576 (1 MB)              | Parquet page size                                                                                                                                                                                                                                  |
 | write.parquet.page-version                          | v1                          | Parquet data page version: v1 (DataPage V1) or v2 (DataPage V2)                                                                                                                                                                                   |
 | write.parquet.page-row-limit                        | 20000                       | Parquet page row limit                                                                                                                                                                                                                             |

parquet/src/main/java/org/apache/iceberg/parquet/Parquet.java

@@ -51,6 +51,8 @@
 import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_CHECK_MIN_RECORD_COUNT_DEFAULT;
 import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES;
 import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT;
+import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED;
+import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED_DEFAULT;
 
 import java.io.File;
 import java.io.IOException;
@@ -372,6 +374,7 @@ public <D> FileAppender<D> build() throws IOException {
       int rowGroupCheckMaxRecordCount = context.rowGroupCheckMaxRecordCount();
       int bloomFilterMaxBytes = context.bloomFilterMaxBytes();
       boolean dictionaryEnabled = context.dictionaryEnabled();
+      boolean trackUncompressedRowGroupSize = context.trackUncompressedRowGroupSize();
 
       if (compressionLevel != null) {
         switch (codec) {
@@ -462,7 +465,8 @@ public <D> FileAppender<D> build() throws IOException {
             parquetProperties,
             metricsConfig,
             writeMode,
-            fileEncryptionProperties);
+            fileEncryptionProperties,
+            trackUncompressedRowGroupSize);
       } else {
         ParquetWriteBuilder<D> parquetWriteBuilder =
             new ParquetWriteBuilder<D>(ParquetIO.file(file))
@@ -510,6 +514,7 @@ static class Context {
       private final Map<String, String> columnBloomFilterEnabled;
       private final Map<String, String> columnStatsEnabled;
       private final boolean dictionaryEnabled;
+      private final boolean trackUncompressedRowGroupSize;
 
       private Context(
           int rowGroupSize,
@@ -526,7 +531,8 @@ private Context(
           Map<String, String> columnBloomFilterNdv,
           Map<String, String> columnBloomFilterEnabled,
           Map<String, String> columnStatsEnabled,
-          boolean dictionaryEnabled) {
+          boolean dictionaryEnabled,
+          boolean trackUncompressedRowGroupSize) {
         this.rowGroupSize = rowGroupSize;
         this.pageSize = pageSize;
         this.pageRowLimit = pageRowLimit;
@@ -542,6 +548,7 @@ private Context(
         this.columnBloomFilterEnabled = columnBloomFilterEnabled;
         this.columnStatsEnabled = columnStatsEnabled;
         this.dictionaryEnabled = dictionaryEnabled;
+        this.trackUncompressedRowGroupSize = trackUncompressedRowGroupSize;
       }
 
       static Context dataContext(Map<String, String> config) {
@@ -615,6 +622,12 @@ static Context dataContext(Map<String, String> config) {
         boolean dictionaryEnabled =
             PropertyUtil.propertyAsBoolean(config, ParquetOutputFormat.ENABLE_DICTIONARY, true);
 
+        boolean trackUncompressedRowGroupSize =
+            PropertyUtil.propertyAsBoolean(
+                config,
+                PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED,
+                PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED_DEFAULT);
+
         return new Context(
             rowGroupSize,
             pageSize,
@@ -630,7 +643,8 @@ static Context dataContext(Map<String, String> config) {
             columnBloomFilterNdv,
             columnBloomFilterEnabled,
             columnStatsEnabled,
-            dictionaryEnabled);
+            dictionaryEnabled,
+            trackUncompressedRowGroupSize);
       }
 
       static Context deleteContext(Map<String, String> config) {
@@ -707,7 +721,8 @@ static Context deleteContext(Map<String, String> config) {
             ImmutableMap.of(),
             ImmutableMap.of(),
             ImmutableMap.of(),
-            dictionaryEnabled);
+            dictionaryEnabled,
+            dataContext.trackUncompressedRowGroupSize());
       }
 
       private static CompressionCodecName toCodec(String codecAsString) {
@@ -786,6 +801,10 @@ Map<String, String> columnStatsEnabled() {
       boolean dictionaryEnabled() {
         return dictionaryEnabled;
       }
+
+      boolean trackUncompressedRowGroupSize() {
+        return trackUncompressedRowGroupSize;
+      }
     }
   }
 

parquet/src/main/java/org/apache/iceberg/parquet/ParquetWriter.java

@@ -59,11 +59,13 @@ class ParquetWriter<T> implements FileAppender<T>, Closeable {
   private final OutputFile output;
   private final Configuration conf;
   private final InternalFileEncryptor fileEncryptor;
+  private final boolean trackUncompressedSize;
 
   private ColumnChunkPageWriteStore pageStore = null;
   private ColumnWriteStore writeStore;
   private long recordCount = 0;
   private long nextCheckRecordCount = 10;
+  private long rowGroupUncompressedSize = 0;
   private boolean closed;
   private ParquetFileWriter writer;
   private int rowGroupOrdinal;
@@ -84,10 +86,12 @@ class ParquetWriter<T> implements FileAppender<T>, Closeable {
       ParquetProperties properties,
       MetricsConfig metricsConfig,
       ParquetFileWriter.Mode writeMode,
-      FileEncryptionProperties encryptionProperties) {
+      FileEncryptionProperties encryptionProperties,
+      boolean trackUncompressedSize) {
     this.schema = schema;
     this.targetRowGroupSize = rowGroupSize;
     this.props = properties;
+    this.trackUncompressedSize = trackUncompressedSize;
     this.metadata = ImmutableMap.copyOf(metadata);
     this.compressor =
         new ParquetCodecFactory(conf, props.getPageSizeThreshold()).getCompressor(codec);
@@ -135,11 +139,21 @@ private void ensureWriterInitialized() {
   @Override
   public void add(T value) {
     recordCount += 1;
-    model.write(0, value);
+    if (trackUncompressedSize) {
+      writeTracked(value);
+    } else {
+      model.write(0, value);
+    }
     writeStore.endRecord();
     checkSize();
   }
 
+  private void writeTracked(T value) {
+    long sizeBefore = writeStore.getBufferedSize();
+    model.write(0, value);
+    rowGroupUncompressedSize += writeStore.getBufferedSize() - sizeBefore;
+  }
+
   @Override
   public Metrics metrics() {
     Preconditions.checkState(closed, "Cannot return metrics for unclosed writer");
@@ -190,21 +204,30 @@ public List<Long> splitOffsets() {
   }
 
   private void checkSize() {
-    if (recordCount >= nextCheckRecordCount) {
-      long bufferedSize = writeStore.getBufferedSize();
-      double avgRecordSize = ((double) bufferedSize) / recordCount;
-
-      if (bufferedSize > (targetRowGroupSize - 2 * avgRecordSize)) {
+    if (trackUncompressedSize) {
+      if (rowGroupUncompressedSize >= targetRowGroupSize) {
         flushRowGroup(false);
-      } else {
-        long remainingSpace = targetRowGroupSize - bufferedSize;
-        long remainingRecords = (long) (remainingSpace / avgRecordSize);
-        this.nextCheckRecordCount =
-            recordCount
-                + Math.min(
-                    Math.max(remainingRecords / 2, props.getMinRowCountForPageSizeCheck()),
-                    props.getMaxRowCountForPageSizeCheck());
+      } else if (recordCount >= nextCheckRecordCount) {
+        evaluateRowGroupSize(rowGroupUncompressedSize, false);
+      }
+    } else if (recordCount >= nextCheckRecordCount) {
+      evaluateRowGroupSize(writeStore.getBufferedSize(), true);
+    }
+  }
+
+  private void evaluateRowGroupSize(long currentSize, boolean useMinRecordCountFloor) {
+    double avgRecordSize = ((double) currentSize) / recordCount;
+    if (currentSize > (targetRowGroupSize - 2 * avgRecordSize)) {
+      flushRowGroup(false);
+    } else {
+      long remainingSpace = targetRowGroupSize - currentSize;
+      long remainingRecords = (long) (remainingSpace / avgRecordSize);
+      long interval = remainingRecords / 2;
+      if (useMinRecordCountFloor) {
+        interval = Math.max(interval, props.getMinRowCountForPageSizeCheck());
       }
+      this.nextCheckRecordCount =
+          recordCount + Math.min(interval, props.getMaxRowCountForPageSizeCheck());
     }
   }
 
@@ -234,6 +257,7 @@ private void startRowGroup() {
             Math.max(recordCount / 2, props.getMinRowCountForPageSizeCheck()),
             props.getMaxRowCountForPageSizeCheck());
     this.recordCount = 0;
+    this.rowGroupUncompressedSize = 0;
 
     this.pageStore =
         new ColumnChunkPageWriteStore(

parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java

@@ -26,6 +26,7 @@
 import java.nio.file.Path;
 import java.util.List;
 import java.util.Optional;
+import java.util.Random;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
 import org.apache.iceberg.DataFile;
@@ -68,6 +69,8 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 public class TestParquetDataWriter {
   private static final Schema SCHEMA =
@@ -543,4 +546,70 @@ protected int resolveColumnIndex(Void engineSchema, String columnName) {
           variantSchema.asStruct(), variantRecords.get(i), writtenRecords.get(i));
     }
   }
+
+  @ParameterizedTest
+  @ValueSource(strings = {"gzip", "snappy", "zstd", "uncompressed"})
+  public void testRowGroupSizeEnforcedWhenCompressionEnabled(String codec) throws IOException {
+    // With uncompressed tracking, row groups split at the configured target
+    DataFile dataFile = writeCompressibleRecords(codec, true);
+
+    assertThat(dataFile.recordCount()).as("Record count should match").isEqualTo(30);
+    assertThat(dataFile.splitOffsets().size())
+        .as("Row group count should reflect enforcement of the target")
+        .isGreaterThanOrEqualTo(3);
+  }
+
+  @Test
+  public void testDefaultPathUsesCompressedSize() throws IOException {
+    // Without uncompressed tracking, compressed bytes never hit the target
+    DataFile dataFile = writeCompressibleRecords("snappy", false);
+    DataFile trackedFile = writeCompressibleRecords("snappy", true);
+
+    assertThat(dataFile.splitOffsets().size())
+        .as("Default path should produce fewer row groups than the tracking path")
+        .isLessThan(trackedFile.splitOffsets().size());
+  }
+
+  // Writes 30 records of 256 KB compressible JSON (~8 MB uncompressed) with a 2 MB target.
+  private DataFile writeCompressibleRecords(String codec, boolean trackUncompressed)
+      throws IOException {
+    OutputFile file = Files.localOutput(createTempFile(temp));
+
+    long targetRowGroupSize = 2 * 1024 * 1024;
+
+    Parquet.DataWriteBuilder builder =
+        Parquet.writeData(file)
+            .schema(SCHEMA)
+            .createWriterFunc(GenericParquetWriter::create)
+            .overwrite()
+            .withSpec(PartitionSpec.unpartitioned())
+            .set("write.parquet.row-group-size-bytes", String.valueOf(targetRowGroupSize))
+            .set("write.parquet.page-size-bytes", "1048576")
+            .set("write.parquet.compression-codec", codec);
+
+    if (trackUncompressed) {
+      builder.set("write.parquet.row-group-size-track-uncompressed", "true");
+    }
+
+    DataWriter<Record> dataWriter = builder.build();
+
+    try (dataWriter) {
+      Random rng = new Random(42);
+      for (int i = 0; i < 30; i++) {
+        GenericRecord record = GenericRecord.create(SCHEMA);
+        record.setField("id", (long) i);
+        StringBuilder sb = new StringBuilder(256 * 1024);
+        sb.append("{\"id\":").append(i).append(",\"values\":[");
+        while (sb.length() < 256 * 1024) {
+          sb.append(rng.nextInt(100000)).append(',');
+        }
+        sb.setCharAt(sb.length() - 1, ']');
+        sb.append('}');
+        record.setField("data", sb.toString());
+        dataWriter.write(record);
+      }
+    }
+
+    return dataWriter.toDataFile();
+  }
 }