huaxingao commented on code in PR #16094:
URL: https://github.com/apache/iceberg/pull/16094#discussion_r4471061834
##########
parquet/src/test/java/org/apache/iceberg/parquet/TestParquet.java:
##########
@@ -774,6 +781,146 @@ private Pair<File, Long> generateFile(
return Pair.of(file, size);
}
+ @Test
+ void globalCompressionCodecAppliesToAllColumns() throws Exception {
+ File file = writeIntAndStringFields(ImmutableMap.of(PARQUET_COMPRESSION,
"snappy"), 5);
+
+ assertThat(readColumnCodecs(file))
+ .containsEntry("int_field", CompressionCodecName.SNAPPY)
+ .containsEntry("string_field", CompressionCodecName.SNAPPY);
+ }
+
+ @Test
+ void perColumnCompressionCodec() throws Exception {
+ File file =
+ writeIntAndStringFields(
+ ImmutableMap.of(
+ PARQUET_COMPRESSION,
+ "zstd",
+ PARQUET_COLUMN_COMPRESSION_CODEC_PREFIX + "int_field",
+ "snappy"),
+ 5);
+
+ assertThat(readColumnCodecs(file))
+ .containsEntry("int_field", CompressionCodecName.SNAPPY)
+ .containsEntry("string_field", CompressionCodecName.ZSTD);
+ }
+
+ @Test
+ void perColumnCompressionCodecForNonExistentColumnIsIgnored() throws
Exception {
+ File file =
+ writeIntAndStringFields(
+ ImmutableMap.of(
+ PARQUET_COMPRESSION,
+ "snappy",
+ PARQUET_COLUMN_COMPRESSION_CODEC_PREFIX + "non_existent_field",
+ "zstd"),
+ 5);
+
+ assertThat(readColumnCodecs(file))
+ .containsEntry("int_field", CompressionCodecName.SNAPPY)
+ .containsEntry("string_field", CompressionCodecName.SNAPPY);
+ }
+
+ @Test
+ void perColumnCompressionLevelIsApplied() throws Exception {
+ File lowLevelFile = writeWithStringFieldZstdLevel(1);
+ File highLevelFile = writeWithStringFieldZstdLevel(19);
+
+ assertThat(readColumnCodecs(highLevelFile))
+ .containsEntry("int_field", CompressionCodecName.SNAPPY)
+ .containsEntry("string_field", CompressionCodecName.ZSTD);
+ assertThat(readColumnSizes(highLevelFile).get("string_field"))
+ .isLessThan(readColumnSizes(lowLevelFile).get("string_field"));
+ }
+
+ private File writeWithStringFieldZstdLevel(int level) throws IOException {
+ Random random = new Random(0);
+ return writeIntAndStringFields(
+ ImmutableMap.of(
+ PARQUET_COMPRESSION,
+ "snappy",
+ PARQUET_COLUMN_COMPRESSION_CODEC_PREFIX + "string_field",
+ "zstd",
+ PARQUET_COLUMN_COMPRESSION_LEVEL_PREFIX + "string_field",
+ String.valueOf(level)),
+ 20_000,
+ i -> "value-" + random.nextInt(5_000) + "-" + random.nextInt(5_000));
+ }
+
+ @Test
+ void invalidCompressionLevelThrows() {
+ assertThatThrownBy(
+ () ->
+ writeIntAndStringFields(
+ ImmutableMap.of(
+ PARQUET_COLUMN_COMPRESSION_CODEC_PREFIX + "int_field",
+ "zstd",
+ PARQUET_COLUMN_COMPRESSION_LEVEL_PREFIX + "int_field",
+ "not-a-number"),
+ 1))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("Invalid compression level for column int_field:
not-a-number");
+ }
+
+ private File writeIntAndStringFields(Map<String, String> properties, int
recordCount)
+ throws IOException {
+ return writeIntAndStringFields(properties, recordCount, i -> "test");
+ }
+
+ private File writeIntAndStringFields(
+ Map<String, String> properties, int recordCount, IntFunction<String>
stringValue)
+ throws IOException {
+ Schema schema =
+ new Schema(
+ optional(1, "int_field", IntegerType.get()),
+ optional(2, "string_field", Types.StringType.get()));
+
+ org.apache.avro.Schema avroSchema =
AvroSchemaUtil.convert(schema.asStruct());
+ List<GenericData.Record> records =
Lists.newArrayListWithCapacity(recordCount);
+ for (int i = 1; i <= recordCount; i++) {
+ GenericData.Record record = new GenericData.Record(avroSchema);
+ record.put("int_field", i);
+ record.put("string_field", stringValue.apply(i));
+ records.add(record);
+ }
+
+ File file = createTempFile(temp);
+ write(
+ file,
+ schema,
+ properties,
+ ParquetAvroWriter::buildWriter,
+ records.toArray(new GenericData.Record[] {}));
+ return file;
+ }
+
+ private Map<String, Long> readColumnSizes(File file) throws IOException {
+ Map<String, Long> sizes = Maps.newHashMap();
+ try (ParquetFileReader reader =
+ ParquetFileReader.open(ParquetIO.file(Files.localInput(file)))) {
+ for (BlockMetaData block : reader.getFooter().getBlocks()) {
+ for (ColumnChunkMetaData column : block.getColumns()) {
+ sizes.merge(column.getPath().toDotString(), column.getTotalSize(),
Long::sum);
+ }
+ }
+ }
+ return sizes;
Review Comment:
Nit: Iceberg style adds an empty line after a control flow block. Could you
add one here before `return sizes;`, in `readColumnCodecs` before `return
codecs;`, and in `Parquet.java` `applyCompressionLevel` after the `if
(compressionLevel == null)` block?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
