Skip to content

Commit e829aaf

Browse files
nssaliannssalian
authored
Parquet: Add opt-in uncompressed row group size tracking (#16327)
* Parquet: Add opt-in uncompressed row group size tracking * Adding test and rename checkSize method * PR comments * test fix * PR comments * Add back the gate outside evaluateRowGroupSize
1 parent 13182ab commit e829aaf

5 files changed

Lines changed: 136 additions & 19 deletions

File tree

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,10 @@ private TableProperties() {}
176176
"write.delete.parquet.row-group-check-max-record-count";
177177
public static final int PARQUET_ROW_GROUP_CHECK_MAX_RECORD_COUNT_DEFAULT = 10000;
178178

179+
public static final String PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED =
180+
"write.parquet.row-group-size-track-uncompressed";
181+
public static final boolean PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED_DEFAULT = false;
182+
179183
public static final String PARQUET_BLOOM_FILTER_MAX_BYTES =
180184
"write.parquet.bloom-filter-max-bytes";
181185
public static final int PARQUET_BLOOM_FILTER_MAX_BYTES_DEFAULT = 1024 * 1024;

‎docs/docs/configuration.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ Iceberg tables support table properties to configure table behavior, like the de
4444
| write.format.default | parquet | Default file format for the table; parquet, avro, or orc |
4545
| write.delete.format.default | data file format | Default delete file format for the table; parquet, avro, or orc |
4646
| write.parquet.row-group-size-bytes | 134217728 (128 MB) | Parquet row group size |
47+
| write.parquet.row-group-size-track-uncompressed | false | Track raw data size to enforce the row group size target accurately with compressing codecs |
4748
| write.parquet.page-size-bytes | 1048576 (1 MB) | Parquet page size |
4849
| write.parquet.page-version | v1 | Parquet data page version: v1 (DataPage V1) or v2 (DataPage V2) |
4950
| write.parquet.page-row-limit | 20000 | Parquet page row limit |

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

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@
5151
import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_CHECK_MIN_RECORD_COUNT_DEFAULT;
5252
import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES;
5353
import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT;
54+
import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED;
55+
import static org.apache.iceberg.TableProperties.PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED_DEFAULT;
5456

5557
import java.io.File;
5658
import java.io.IOException;
@@ -372,6 +374,7 @@ public <D> FileAppender<D> build() throws IOException {
372374
int rowGroupCheckMaxRecordCount = context.rowGroupCheckMaxRecordCount();
373375
int bloomFilterMaxBytes = context.bloomFilterMaxBytes();
374376
boolean dictionaryEnabled = context.dictionaryEnabled();
377+
boolean trackUncompressedRowGroupSize = context.trackUncompressedRowGroupSize();
375378

376379
if (compressionLevel != null) {
377380
switch (codec) {
@@ -462,7 +465,8 @@ public <D> FileAppender<D> build() throws IOException {
462465
parquetProperties,
463466
metricsConfig,
464467
writeMode,
465-
fileEncryptionProperties);
468+
fileEncryptionProperties,
469+
trackUncompressedRowGroupSize);
466470
} else {
467471
ParquetWriteBuilder<D> parquetWriteBuilder =
468472
new ParquetWriteBuilder<D>(ParquetIO.file(file))
@@ -510,6 +514,7 @@ static class Context {
510514
private final Map<String, String> columnBloomFilterEnabled;
511515
private final Map<String, String> columnStatsEnabled;
512516
private final boolean dictionaryEnabled;
517+
private final boolean trackUncompressedRowGroupSize;
513518

514519
private Context(
515520
int rowGroupSize,
@@ -526,7 +531,8 @@ private Context(
526531
Map<String, String> columnBloomFilterNdv,
527532
Map<String, String> columnBloomFilterEnabled,
528533
Map<String, String> columnStatsEnabled,
529-
boolean dictionaryEnabled) {
534+
boolean dictionaryEnabled,
535+
boolean trackUncompressedRowGroupSize) {
530536
this.rowGroupSize = rowGroupSize;
531537
this.pageSize = pageSize;
532538
this.pageRowLimit = pageRowLimit;
@@ -542,6 +548,7 @@ private Context(
542548
this.columnBloomFilterEnabled = columnBloomFilterEnabled;
543549
this.columnStatsEnabled = columnStatsEnabled;
544550
this.dictionaryEnabled = dictionaryEnabled;
551+
this.trackUncompressedRowGroupSize = trackUncompressedRowGroupSize;
545552
}
546553

547554
static Context dataContext(Map<String, String> config) {
@@ -615,6 +622,12 @@ static Context dataContext(Map<String, String> config) {
615622
boolean dictionaryEnabled =
616623
PropertyUtil.propertyAsBoolean(config, ParquetOutputFormat.ENABLE_DICTIONARY, true);
617624

625+
boolean trackUncompressedRowGroupSize =
626+
PropertyUtil.propertyAsBoolean(
627+
config,
628+
PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED,
629+
PARQUET_ROW_GROUP_SIZE_TRACK_UNCOMPRESSED_DEFAULT);
630+
618631
return new Context(
619632
rowGroupSize,
620633
pageSize,
@@ -630,7 +643,8 @@ static Context dataContext(Map<String, String> config) {
630643
columnBloomFilterNdv,
631644
columnBloomFilterEnabled,
632645
columnStatsEnabled,
633-
dictionaryEnabled);
646+
dictionaryEnabled,
647+
trackUncompressedRowGroupSize);
634648
}
635649

636650
static Context deleteContext(Map<String, String> config) {
@@ -707,7 +721,8 @@ static Context deleteContext(Map<String, String> config) {
707721
ImmutableMap.of(),
708722
ImmutableMap.of(),
709723
ImmutableMap.of(),
710-
dictionaryEnabled);
724+
dictionaryEnabled,
725+
dataContext.trackUncompressedRowGroupSize());
711726
}
712727

713728
private static CompressionCodecName toCodec(String codecAsString) {
@@ -786,6 +801,10 @@ Map<String, String> columnStatsEnabled() {
786801
boolean dictionaryEnabled() {
787802
return dictionaryEnabled;
788803
}
804+
805+
boolean trackUncompressedRowGroupSize() {
806+
return trackUncompressedRowGroupSize;
807+
}
789808
}
790809
}
791810

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

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -59,11 +59,13 @@ class ParquetWriter<T> implements FileAppender<T>, Closeable {
5959
private final OutputFile output;
6060
private final Configuration conf;
6161
private final InternalFileEncryptor fileEncryptor;
62+
private final boolean trackUncompressedSize;
6263

6364
private ColumnChunkPageWriteStore pageStore = null;
6465
private ColumnWriteStore writeStore;
6566
private long recordCount = 0;
6667
private long nextCheckRecordCount = 10;
68+
private long rowGroupUncompressedSize = 0;
6769
private boolean closed;
6870
private ParquetFileWriter writer;
6971
private int rowGroupOrdinal;
@@ -84,10 +86,12 @@ class ParquetWriter<T> implements FileAppender<T>, Closeable {
8486
ParquetProperties properties,
8587
MetricsConfig metricsConfig,
8688
ParquetFileWriter.Mode writeMode,
87-
FileEncryptionProperties encryptionProperties) {
89+
FileEncryptionProperties encryptionProperties,
90+
boolean trackUncompressedSize) {
8891
this.schema = schema;
8992
this.targetRowGroupSize = rowGroupSize;
9093
this.props = properties;
94+
this.trackUncompressedSize = trackUncompressedSize;
9195
this.metadata = ImmutableMap.copyOf(metadata);
9296
this.compressor =
9397
new ParquetCodecFactory(conf, props.getPageSizeThreshold()).getCompressor(codec);
@@ -135,11 +139,21 @@ private void ensureWriterInitialized() {
135139
@Override
136140
public void add(T value) {
137141
recordCount += 1;
138-
model.write(0, value);
142+
if (trackUncompressedSize) {
143+
writeTracked(value);
144+
} else {
145+
model.write(0, value);
146+
}
139147
writeStore.endRecord();
140148
checkSize();
141149
}
142150

151+
private void writeTracked(T value) {
152+
long sizeBefore = writeStore.getBufferedSize();
153+
model.write(0, value);
154+
rowGroupUncompressedSize += writeStore.getBufferedSize() - sizeBefore;
155+
}
156+
143157
@Override
144158
public Metrics metrics() {
145159
Preconditions.checkState(closed, "Cannot return metrics for unclosed writer");
@@ -190,21 +204,30 @@ public List<Long> splitOffsets() {
190204
}
191205

192206
private void checkSize() {
193-
if (recordCount >= nextCheckRecordCount) {
194-
long bufferedSize = writeStore.getBufferedSize();
195-
double avgRecordSize = ((double) bufferedSize) / recordCount;
196-
197-
if (bufferedSize > (targetRowGroupSize - 2 * avgRecordSize)) {
207+
if (trackUncompressedSize) {
208+
if (rowGroupUncompressedSize >= targetRowGroupSize) {
198209
flushRowGroup(false);
199-
} else {
200-
long remainingSpace = targetRowGroupSize - bufferedSize;
201-
long remainingRecords = (long) (remainingSpace / avgRecordSize);
202-
this.nextCheckRecordCount =
203-
recordCount
204-
+ Math.min(
205-
Math.max(remainingRecords / 2, props.getMinRowCountForPageSizeCheck()),
206-
props.getMaxRowCountForPageSizeCheck());
210+
} else if (recordCount >= nextCheckRecordCount) {
211+
evaluateRowGroupSize(rowGroupUncompressedSize, false);
212+
}
213+
} else if (recordCount >= nextCheckRecordCount) {
214+
evaluateRowGroupSize(writeStore.getBufferedSize(), true);
215+
}
216+
}
217+
218+
private void evaluateRowGroupSize(long currentSize, boolean useMinRecordCountFloor) {
219+
double avgRecordSize = ((double) currentSize) / recordCount;
220+
if (currentSize > (targetRowGroupSize - 2 * avgRecordSize)) {
221+
flushRowGroup(false);
222+
} else {
223+
long remainingSpace = targetRowGroupSize - currentSize;
224+
long remainingRecords = (long) (remainingSpace / avgRecordSize);
225+
long interval = remainingRecords / 2;
226+
if (useMinRecordCountFloor) {
227+
interval = Math.max(interval, props.getMinRowCountForPageSizeCheck());
207228
}
229+
this.nextCheckRecordCount =
230+
recordCount + Math.min(interval, props.getMaxRowCountForPageSizeCheck());
208231
}
209232
}
210233

@@ -234,6 +257,7 @@ private void startRowGroup() {
234257
Math.max(recordCount / 2, props.getMinRowCountForPageSizeCheck()),
235258
props.getMaxRowCountForPageSizeCheck());
236259
this.recordCount = 0;
260+
this.rowGroupUncompressedSize = 0;
237261

238262
this.pageStore =
239263
new ColumnChunkPageWriteStore(

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import java.nio.file.Path;
2727
import java.util.List;
2828
import java.util.Optional;
29+
import java.util.Random;
2930
import java.util.function.Function;
3031
import java.util.function.UnaryOperator;
3132
import org.apache.iceberg.DataFile;
@@ -68,6 +69,8 @@
6869
import org.junit.jupiter.api.BeforeEach;
6970
import org.junit.jupiter.api.Test;
7071
import org.junit.jupiter.api.io.TempDir;
72+
import org.junit.jupiter.params.ParameterizedTest;
73+
import org.junit.jupiter.params.provider.ValueSource;
7174

7275
public class TestParquetDataWriter {
7376
private static final Schema SCHEMA =
@@ -543,4 +546,70 @@ protected int resolveColumnIndex(Void engineSchema, String columnName) {
543546
variantSchema.asStruct(), variantRecords.get(i), writtenRecords.get(i));
544547
}
545548
}
549+
550+
@ParameterizedTest
551+
@ValueSource(strings = {"gzip", "snappy", "zstd", "uncompressed"})
552+
public void testRowGroupSizeEnforcedWhenCompressionEnabled(String codec) throws IOException {
553+
// With uncompressed tracking, row groups split at the configured target
554+
DataFile dataFile = writeCompressibleRecords(codec, true);
555+
556+
assertThat(dataFile.recordCount()).as("Record count should match").isEqualTo(30);
557+
assertThat(dataFile.splitOffsets().size())
558+
.as("Row group count should reflect enforcement of the target")
559+
.isGreaterThanOrEqualTo(3);
560+
}
561+
562+
@Test
563+
public void testDefaultPathUsesCompressedSize() throws IOException {
564+
// Without uncompressed tracking, compressed bytes never hit the target
565+
DataFile dataFile = writeCompressibleRecords("snappy", false);
566+
DataFile trackedFile = writeCompressibleRecords("snappy", true);
567+
568+
assertThat(dataFile.splitOffsets().size())
569+
.as("Default path should produce fewer row groups than the tracking path")
570+
.isLessThan(trackedFile.splitOffsets().size());
571+
}
572+
573+
// Writes 30 records of 256 KB compressible JSON (~8 MB uncompressed) with a 2 MB target.
574+
private DataFile writeCompressibleRecords(String codec, boolean trackUncompressed)
575+
throws IOException {
576+
OutputFile file = Files.localOutput(createTempFile(temp));
577+
578+
long targetRowGroupSize = 2 * 1024 * 1024;
579+
580+
Parquet.DataWriteBuilder builder =
581+
Parquet.writeData(file)
582+
.schema(SCHEMA)
583+
.createWriterFunc(GenericParquetWriter::create)
584+
.overwrite()
585+
.withSpec(PartitionSpec.unpartitioned())
586+
.set("write.parquet.row-group-size-bytes", String.valueOf(targetRowGroupSize))
587+
.set("write.parquet.page-size-bytes", "1048576")
588+
.set("write.parquet.compression-codec", codec);
589+
590+
if (trackUncompressed) {
591+
builder.set("write.parquet.row-group-size-track-uncompressed", "true");
592+
}
593+
594+
DataWriter<Record> dataWriter = builder.build();
595+
596+
try (dataWriter) {
597+
Random rng = new Random(42);
598+
for (int i = 0; i < 30; i++) {
599+
GenericRecord record = GenericRecord.create(SCHEMA);
600+
record.setField("id", (long) i);
601+
StringBuilder sb = new StringBuilder(256 * 1024);
602+
sb.append("{\"id\":").append(i).append(",\"values\":[");
603+
while (sb.length() < 256 * 1024) {
604+
sb.append(rng.nextInt(100000)).append(',');
605+
}
606+
sb.setCharAt(sb.length() - 1, ']');
607+
sb.append('}');
608+
record.setField("data", sb.toString());
609+
dataWriter.write(record);
610+
}
611+
}
612+
613+
return dataWriter.toDataFile();
614+
}
546615
}

0 commit comments

Comments
 (0)