From cfa3f283795a4cc8277d9967fc8bd235ddbe0cc3 Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Sat, 1 Aug 2026 18:59:38 -0700 Subject: [PATCH 1/3] Parquet: Add round-trip test for encrypted shredded variant columns --- .../parquet/TestParquetDataWriter.java | 129 ++++++++++++++++++ 1 file changed, 129 insertions(+) diff --git a/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java index 00891b507eef..7e58e78c69b3 100644 --- a/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java +++ b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java @@ -20,10 +20,12 @@ import static org.apache.iceberg.parquet.ParquetWritingTestUtils.createTempFile; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.file.Path; +import java.security.SecureRandom; import java.util.List; import java.util.Optional; import java.util.Random; @@ -61,6 +63,9 @@ import org.apache.iceberg.variants.VariantTestUtil; import org.apache.iceberg.variants.VariantValue; import org.apache.iceberg.variants.Variants; +import org.apache.parquet.ParquetReadOptions; +import org.apache.parquet.crypto.FileDecryptionProperties; +import org.apache.parquet.crypto.ParquetCryptoRuntimeException; import org.apache.parquet.example.data.Group; import org.apache.parquet.hadoop.ParquetFileReader; import org.apache.parquet.hadoop.ParquetReader; @@ -607,6 +612,130 @@ protected int resolveColumnIndex(Void engineSchema, String columnName) { } } + @Test + public void testFormatModelVariantShreddingWithEncryption() throws IOException { + Schema variantSchema = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "v", Types.VariantType.get())); + + VariantShreddingAnalyzer analyzer = + new VariantShreddingAnalyzer() { + @Override + protected List extractVariantValues(List rows, int idx) { + List values = Lists.newArrayList(); + for (Record row : rows) { + Object obj = row.get(idx); + if (obj instanceof Variant) { + values.add(((Variant) obj).value()); + } + } + return values; + } + + @Override + protected int resolveColumnIndex(Void engineSchema, String columnName) { + return variantSchema.columns().indexOf(variantSchema.findField(columnName)); + } + }; + + ByteBuffer metadataBuffer = VariantTestUtil.createMetadata(ImmutableList.of("a", "b"), true); + VariantMetadata metadata = Variants.metadata(metadataBuffer); + ByteBuffer objectBuffer = + VariantTestUtil.createObject( + metadataBuffer, + ImmutableMap.of("a", Variants.of(123456789), "b", Variants.of("string"))); + Variant variant = Variant.of(metadata, Variants.value(metadata, objectBuffer)); + + GenericRecord record = GenericRecord.create(variantSchema); + List variantRecords = + ImmutableList.of( + record.copy(ImmutableMap.of("id", 1L, "v", variant)), + record.copy(ImmutableMap.of("id", 2L, "v", variant)), + record.copy(ImmutableMap.of("id", 3L, "v", variant))); + + ParquetFormatModel> model = + ParquetFormatModel.create( + Record.class, + Void.class, + (icebergSchema, messageType, engineSchema) -> + GenericParquetWriter.create(icebergSchema, messageType), + (icebergSchema, fileSchema, engineSchema, idToConstant) -> + GenericParquetReaders.buildReader(icebergSchema, fileSchema), + analyzer, + (Function>) unused -> input -> input); + + OutputFile encryptedFile = Files.localOutput(createTempFile(temp)); + ByteBuffer fileDek = ByteBuffer.allocate(16); + ByteBuffer aadPrefix = ByteBuffer.allocate(16); + SecureRandom random = new SecureRandom(); + random.nextBytes(fileDek.array()); + random.nextBytes(aadPrefix.array()); + + try (FileAppender appender = + model + .writeBuilder(EncryptedFiles.plainAsEncryptedOutput(encryptedFile)) + .schema(variantSchema) + .withFileEncryptionKey(fileDek) + .withAADPrefix(aadPrefix) + .setAll( + ImmutableMap.of( + TableProperties.PARQUET_SHRED_VARIANTS, "true", + TableProperties.PARQUET_VARIANT_BUFFER_SIZE, "2")) + .content(FileContent.DATA) + .build()) { + assertThat(appender).isInstanceOf(BufferedFileAppender.class); + appender.addAll(variantRecords); + } + + assertThatThrownBy( + () -> + Parquet.read(encryptedFile.toInputFile()) + .project(variantSchema) + .createReaderFunc( + fileSchema -> GenericParquetReaders.buildReader(variantSchema, fileSchema)) + .build() + .iterator()) + .isInstanceOf(ParquetCryptoRuntimeException.class) + .hasMessage("Trying to read file with encrypted footer. No keys available"); + + List writtenRecords; + try (CloseableIterable reader = + Parquet.read(encryptedFile.toInputFile()) + .project(variantSchema) + .withFileEncryptionKey(fileDek) + .withAADPrefix(aadPrefix) + .createReaderFunc( + fileSchema -> GenericParquetReaders.buildReader(variantSchema, fileSchema)) + .build()) { + writtenRecords = Lists.newArrayList(reader); + } + + assertThat(writtenRecords).hasSameSizeAs(variantRecords); + for (int i = 0; i < variantRecords.size(); i++) { + InternalTestHelpers.assertEquals( + variantSchema.asStruct(), variantRecords.get(i), writtenRecords.get(i)); + } + + try (ParquetFileReader fileReader = + ParquetFileReader.open( + ParquetIO.file(encryptedFile.toInputFile()), + ParquetReadOptions.builder() + .withDecryption( + FileDecryptionProperties.builder() + .withFooterKey(fileDek.array()) + .withAADPrefix(aadPrefix.array()) + .build()) + .build())) { + GroupType variantType = + fileReader.getFooter().getFileMetaData().getSchema().getType("v").asGroupType(); + assertThat(variantType.containsField("typed_value")).isTrue(); + GroupType typedValue = variantType.getType("typed_value").asGroupType(); + assertThat(typedValue.containsField("a")).isTrue(); + assertThat(typedValue.containsField("b")).isTrue(); + } + } + @ParameterizedTest @ValueSource(strings = {"gzip", "snappy", "zstd", "uncompressed"}) public void testRowGroupSizeEnforcedWhenCompressionEnabled(String codec) throws IOException { From 9deb2ef665b46b14bb57ff88b08a5741c0a8498b Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Sat, 1 Aug 2026 23:53:59 -0700 Subject: [PATCH 2/3] PR comment fix --- .../parquet/TestParquetDataWriter.java | 153 +++++++----------- 1 file changed, 58 insertions(+), 95 deletions(-) diff --git a/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java index 7e58e78c69b3..04e7f2398591 100644 --- a/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java +++ b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java @@ -63,7 +63,6 @@ import org.apache.iceberg.variants.VariantTestUtil; import org.apache.iceberg.variants.VariantValue; import org.apache.iceberg.variants.Variants; -import org.apache.parquet.ParquetReadOptions; import org.apache.parquet.crypto.FileDecryptionProperties; import org.apache.parquet.crypto.ParquetCryptoRuntimeException; import org.apache.parquet.example.data.Group; @@ -85,6 +84,11 @@ public class TestParquetDataWriter { Types.NestedField.optional(2, "data", Types.StringType.get()), Types.NestedField.optional(3, "binary", Types.BinaryType.get())); + private static final Schema VARIANT_SHREDDING_SCHEMA = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "v", Types.VariantType.get())); + private List records; @TempDir private Path temp; @@ -485,13 +489,8 @@ public void testWriteBuilderReturnsDirectAppenderWithNullAnalyzer() throws IOExc } } - @Test - public void testFormatModelVariantShreddingRoundTrip() throws IOException { - Schema variantSchema = - new Schema( - Types.NestedField.required(1, "id", Types.LongType.get()), - Types.NestedField.optional(2, "v", Types.VariantType.get())); - + private static ParquetFormatModel> variantShreddingModel( + Schema schema) { VariantShreddingAnalyzer analyzer = new VariantShreddingAnalyzer() { @Override @@ -508,40 +507,43 @@ protected List extractVariantValues(List rows, int idx) { @Override protected int resolveColumnIndex(Void engineSchema, String columnName) { - // GenericRecord uses schema column order - return variantSchema.columns().indexOf(variantSchema.findField(columnName)); + return schema.columns().indexOf(schema.findField(columnName)); } }; + return ParquetFormatModel.create( + Record.class, + Void.class, + (icebergSchema, messageType, engineSchema) -> + GenericParquetWriter.create(icebergSchema, messageType), + (icebergSchema, fileSchema, engineSchema, idToConstant) -> + GenericParquetReaders.buildReader(icebergSchema, fileSchema), + analyzer, + (Function>) unused -> input -> input); + } + + private static List variantShreddingRecords(Schema schema, int aValue, String bValue) { ByteBuffer metadataBuffer = VariantTestUtil.createMetadata(ImmutableList.of("a", "b"), true); VariantMetadata metadata = Variants.metadata(metadataBuffer); ByteBuffer objectBuffer = VariantTestUtil.createObject( - metadataBuffer, - ImmutableMap.of( - "a", Variants.of(42), - "b", Variants.of("hello"))); + metadataBuffer, ImmutableMap.of("a", Variants.of(aValue), "b", Variants.of(bValue))); Variant variant = Variant.of(metadata, Variants.value(metadata, objectBuffer)); - GenericRecord record = GenericRecord.create(variantSchema); - List variantRecords = - ImmutableList.of( - record.copy(ImmutableMap.of("id", 1L, "v", variant)), - record.copy(ImmutableMap.of("id", 2L, "v", variant)), - record.copy(ImmutableMap.of("id", 3L, "v", variant))); + GenericRecord record = GenericRecord.create(schema); + return ImmutableList.of( + record.copy(ImmutableMap.of("id", 1L, "v", variant)), + record.copy(ImmutableMap.of("id", 2L, "v", variant)), + record.copy(ImmutableMap.of("id", 3L, "v", variant))); + } + @Test + public void testFormatModelVariantShreddingRoundTrip() throws IOException { + Schema variantSchema = VARIANT_SHREDDING_SCHEMA; + List variantRecords = variantShreddingRecords(variantSchema, 42, "hello"); OutputFile outputFile = Files.localOutput(createTempFile(temp)); - ParquetFormatModel> model = - ParquetFormatModel.create( - Record.class, - Void.class, - (icebergSchema, messageType, engineSchema) -> - GenericParquetWriter.create(icebergSchema, messageType), - (icebergSchema, fileSchema, engineSchema, idToConstant) -> - GenericParquetReaders.buildReader(icebergSchema, fileSchema), - analyzer, - (Function>) unused -> record1 -> record1); + variantShreddingModel(variantSchema); try (FileAppender appender = model @@ -614,56 +616,10 @@ protected int resolveColumnIndex(Void engineSchema, String columnName) { @Test public void testFormatModelVariantShreddingWithEncryption() throws IOException { - Schema variantSchema = - new Schema( - Types.NestedField.required(1, "id", Types.LongType.get()), - Types.NestedField.optional(2, "v", Types.VariantType.get())); - - VariantShreddingAnalyzer analyzer = - new VariantShreddingAnalyzer() { - @Override - protected List extractVariantValues(List rows, int idx) { - List values = Lists.newArrayList(); - for (Record row : rows) { - Object obj = row.get(idx); - if (obj instanceof Variant) { - values.add(((Variant) obj).value()); - } - } - return values; - } - - @Override - protected int resolveColumnIndex(Void engineSchema, String columnName) { - return variantSchema.columns().indexOf(variantSchema.findField(columnName)); - } - }; - - ByteBuffer metadataBuffer = VariantTestUtil.createMetadata(ImmutableList.of("a", "b"), true); - VariantMetadata metadata = Variants.metadata(metadataBuffer); - ByteBuffer objectBuffer = - VariantTestUtil.createObject( - metadataBuffer, - ImmutableMap.of("a", Variants.of(123456789), "b", Variants.of("string"))); - Variant variant = Variant.of(metadata, Variants.value(metadata, objectBuffer)); - - GenericRecord record = GenericRecord.create(variantSchema); - List variantRecords = - ImmutableList.of( - record.copy(ImmutableMap.of("id", 1L, "v", variant)), - record.copy(ImmutableMap.of("id", 2L, "v", variant)), - record.copy(ImmutableMap.of("id", 3L, "v", variant))); - + Schema variantSchema = VARIANT_SHREDDING_SCHEMA; + List variantRecords = variantShreddingRecords(variantSchema, 123456789, "string"); ParquetFormatModel> model = - ParquetFormatModel.create( - Record.class, - Void.class, - (icebergSchema, messageType, engineSchema) -> - GenericParquetWriter.create(icebergSchema, messageType), - (icebergSchema, fileSchema, engineSchema, idToConstant) -> - GenericParquetReaders.buildReader(icebergSchema, fileSchema), - analyzer, - (Function>) unused -> input -> input); + variantShreddingModel(variantSchema); OutputFile encryptedFile = Files.localOutput(createTempFile(temp)); ByteBuffer fileDek = ByteBuffer.allocate(16); @@ -717,22 +673,29 @@ protected int resolveColumnIndex(Void engineSchema, String columnName) { variantSchema.asStruct(), variantRecords.get(i), writtenRecords.get(i)); } - try (ParquetFileReader fileReader = - ParquetFileReader.open( - ParquetIO.file(encryptedFile.toInputFile()), - ParquetReadOptions.builder() - .withDecryption( - FileDecryptionProperties.builder() - .withFooterKey(fileDek.array()) - .withAADPrefix(aadPrefix.array()) - .build()) - .build())) { - GroupType variantType = - fileReader.getFooter().getFileMetaData().getSchema().getType("v").asGroupType(); - assertThat(variantType.containsField("typed_value")).isTrue(); - GroupType typedValue = variantType.getType("typed_value").asGroupType(); - assertThat(typedValue.containsField("a")).isTrue(); - assertThat(typedValue.containsField("b")).isTrue(); + try (ParquetReader rawReader = + ParquetReader.builder( + new GroupReadSupport(), new org.apache.hadoop.fs.Path(encryptedFile.location())) + .withDecryption( + FileDecryptionProperties.builder() + .withFooterKey(fileDek.array()) + .withAADPrefix(aadPrefix.array()) + .build()) + .build()) { + Group row = rawReader.read(); + Group variantData = row.getGroup("v", 0); + + assertThat(variantData.getFieldRepetitionCount("value")) + .as("value should be absent when fully shredded") + .isEqualTo(0); + + Group typedValue = variantData.getGroup("typed_value", 0); + assertThat(typedValue.getGroup("a", 0).getInteger("typed_value", 0)) + .as("typed_value.a should contain 123456789") + .isEqualTo(123456789); + assertThat(typedValue.getGroup("b", 0).getString("typed_value", 0)) + .as("typed_value.b should contain string") + .isEqualTo("string"); } } From f70e4e2d7d3159834089e0403ebc707ad56b3629 Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Sun, 2 Aug 2026 10:52:06 -0700 Subject: [PATCH 3/3] Dedupe the tests --- .../parquet/TestParquetDataWriter.java | 344 ++++++++---------- 1 file changed, 150 insertions(+), 194 deletions(-) diff --git a/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java index 04e7f2398591..ae7d5441c91c 100644 --- a/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java +++ b/parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java @@ -59,10 +59,10 @@ import org.apache.iceberg.types.Types; import org.apache.iceberg.util.RandomUtil; import org.apache.iceberg.variants.Variant; -import org.apache.iceberg.variants.VariantMetadata; import org.apache.iceberg.variants.VariantTestUtil; import org.apache.iceberg.variants.VariantValue; import org.apache.iceberg.variants.Variants; +import org.apache.parquet.ParquetReadOptions; import org.apache.parquet.crypto.FileDecryptionProperties; import org.apache.parquet.crypto.ParquetCryptoRuntimeException; import org.apache.parquet.example.data.Group; @@ -389,17 +389,9 @@ public void testDataWriterWithVariantShredding() throws IOException { .add(Types.NestedField.optional(4, "variant", Types.VariantType.get())) .build()); - ByteBuffer metadataBuffer = VariantTestUtil.createMetadata(ImmutableList.of("a", "b"), true); - VariantMetadata metadata = Variants.metadata(metadataBuffer); - - ByteBuffer objectBuffer = - VariantTestUtil.createObject( - metadataBuffer, - ImmutableMap.of( - "a", Variants.of(123456789), - "b", Variants.of("string"))); - - Variant variant = Variant.of(metadata, Variants.value(metadata, objectBuffer)); + Variant variant = + VariantTestUtil.variant( + ImmutableMap.of("a", Variants.of(123456789), "b", Variants.of("string"))); // Create records with variant data GenericRecord record = GenericRecord.create(variantSchema); @@ -415,11 +407,6 @@ public void testDataWriterWithVariantShredding() throws IOException { @Test public void testShreddingWriteReturnsBufferedAppender() throws IOException { - Schema variantSchema = - new Schema( - Types.NestedField.required(1, "id", Types.LongType.get()), - Types.NestedField.optional(2, "v", Types.VariantType.get())); - VariantShreddingAnalyzer testAnalyzer = new VariantShreddingAnalyzer() { @Override @@ -436,20 +423,12 @@ protected int resolveColumnIndex(Void engineSchema, String columnName) { OutputFile outputFile = Files.localOutput(createTempFile(temp)); ParquetFormatModel> model = - ParquetFormatModel.create( - Record.class, - Void.class, - (icebergSchema, messageType, engineSchema) -> - GenericParquetWriter.create(icebergSchema, messageType), - (icebergSchema, fileSchema, engineSchema, idToConstant) -> - GenericParquetReaders.buildReader(icebergSchema, fileSchema), - testAnalyzer, - (Function>) unused -> record -> record); + variantModel(testAnalyzer, unused -> record -> record); try (FileAppender appender = model .writeBuilder(EncryptedFiles.plainAsEncryptedOutput(outputFile)) - .schema(variantSchema) + .schema(VARIANT_SHREDDING_SCHEMA) .setAll(ImmutableMap.of(TableProperties.PARQUET_SHRED_VARIANTS, "true")) .content(FileContent.DATA) .build()) { @@ -459,28 +438,14 @@ protected int resolveColumnIndex(Void engineSchema, String columnName) { @Test public void testWriteBuilderReturnsDirectAppenderWithNullAnalyzer() throws IOException { - Schema variantSchema = - new Schema( - Types.NestedField.required(1, "id", Types.LongType.get()), - Types.NestedField.optional(2, "v", Types.VariantType.get())); - OutputFile outputFile = Files.localOutput(createTempFile(temp)); - ParquetFormatModel> model = - ParquetFormatModel.create( - Record.class, - Void.class, - (icebergSchema, messageType, engineSchema) -> - GenericParquetWriter.create(icebergSchema, messageType), - (icebergSchema, fileSchema, engineSchema, idToConstant) -> - GenericParquetReaders.buildReader(icebergSchema, fileSchema), - null, - (Function>) null); + ParquetFormatModel> model = variantModel(null, null); try (FileAppender appender = model .writeBuilder(EncryptedFiles.plainAsEncryptedOutput(outputFile)) - .schema(variantSchema) + .schema(VARIANT_SHREDDING_SCHEMA) .setAll(ImmutableMap.of(TableProperties.PARQUET_SHRED_VARIANTS, "true")) .content(FileContent.DATA) .build()) { @@ -489,66 +454,16 @@ public void testWriteBuilderReturnsDirectAppenderWithNullAnalyzer() throws IOExc } } - private static ParquetFormatModel> variantShreddingModel( - Schema schema) { - VariantShreddingAnalyzer analyzer = - new VariantShreddingAnalyzer() { - @Override - protected List extractVariantValues(List rows, int idx) { - List values = Lists.newArrayList(); - for (Record row : rows) { - Object obj = row.get(idx); - if (obj instanceof Variant) { - values.add(((Variant) obj).value()); - } - } - return values; - } - - @Override - protected int resolveColumnIndex(Void engineSchema, String columnName) { - return schema.columns().indexOf(schema.findField(columnName)); - } - }; - - return ParquetFormatModel.create( - Record.class, - Void.class, - (icebergSchema, messageType, engineSchema) -> - GenericParquetWriter.create(icebergSchema, messageType), - (icebergSchema, fileSchema, engineSchema, idToConstant) -> - GenericParquetReaders.buildReader(icebergSchema, fileSchema), - analyzer, - (Function>) unused -> input -> input); - } - - private static List variantShreddingRecords(Schema schema, int aValue, String bValue) { - ByteBuffer metadataBuffer = VariantTestUtil.createMetadata(ImmutableList.of("a", "b"), true); - VariantMetadata metadata = Variants.metadata(metadataBuffer); - ByteBuffer objectBuffer = - VariantTestUtil.createObject( - metadataBuffer, ImmutableMap.of("a", Variants.of(aValue), "b", Variants.of(bValue))); - Variant variant = Variant.of(metadata, Variants.value(metadata, objectBuffer)); - - GenericRecord record = GenericRecord.create(schema); - return ImmutableList.of( - record.copy(ImmutableMap.of("id", 1L, "v", variant)), - record.copy(ImmutableMap.of("id", 2L, "v", variant)), - record.copy(ImmutableMap.of("id", 3L, "v", variant))); - } - @Test public void testFormatModelVariantShreddingRoundTrip() throws IOException { - Schema variantSchema = VARIANT_SHREDDING_SCHEMA; - List variantRecords = variantShreddingRecords(variantSchema, 42, "hello"); + List variantRecords = variantShreddingRecords(42, "hello"); OutputFile outputFile = Files.localOutput(createTempFile(temp)); - ParquetFormatModel> model = - variantShreddingModel(variantSchema); + ParquetFormatModel> model = variantShreddingModel(); try (FileAppender appender = model .writeBuilder(EncryptedFiles.plainAsEncryptedOutput(outputFile)) - .schema(variantSchema) + .schema(VARIANT_SHREDDING_SCHEMA) .setAll( ImmutableMap.of( TableProperties.PARQUET_SHRED_VARIANTS, "true", @@ -561,65 +476,15 @@ public void testFormatModelVariantShreddingRoundTrip() throws IOException { } } - // Verify shredded Parquet schema - try (ParquetFileReader reader = - ParquetFileReader.open(ParquetIO.file(outputFile.toInputFile()))) { - MessageType parquetSchema = reader.getFooter().getFileMetaData().getSchema(); - GroupType variantGroup = parquetSchema.getType("v").asGroupType(); - assertThat(variantGroup.containsField("metadata")).isTrue(); - assertThat(variantGroup.containsField("value")).isTrue(); - assertThat(variantGroup.containsField("typed_value")).isTrue(); - - GroupType typedValue = variantGroup.getType("typed_value").asGroupType(); - assertThat(typedValue.containsField("a")).isTrue(); - assertThat(typedValue.containsField("b")).isTrue(); - } + assertShreddedVariant(outputFile, null, 42, "hello"); - // Verify data is in typed columns by reading raw Parquet groups - try (ParquetReader rawReader = - ParquetReader.builder( - new GroupReadSupport(), new org.apache.hadoop.fs.Path(outputFile.location())) - .build()) { - Group row = rawReader.read(); - Group variantData = row.getGroup("v", 0); - - assertThat(variantData.getFieldRepetitionCount("value")) - .as("value should be absent when fully shredded") - .isEqualTo(0); - - Group typedValue = variantData.getGroup("typed_value", 0); - assertThat(typedValue.getGroup("a", 0).getInteger("typed_value", 0)) - .as("typed_value.a should contain 42") - .isEqualTo(42); - assertThat(typedValue.getGroup("b", 0).getString("typed_value", 0)) - .as("typed_value.b should contain hello") - .isEqualTo("hello"); - } - - // Verify data round-trips - List writtenRecords; - try (CloseableIterable reader = - Parquet.read(outputFile.toInputFile()) - .project(variantSchema) - .createReaderFunc( - fileSchema -> GenericParquetReaders.buildReader(variantSchema, fileSchema)) - .build()) { - writtenRecords = Lists.newArrayList(reader); - } - - assertThat(writtenRecords).hasSameSizeAs(variantRecords); - for (int i = 0; i < variantRecords.size(); i++) { - InternalTestHelpers.assertEquals( - variantSchema.asStruct(), variantRecords.get(i), writtenRecords.get(i)); - } + assertRoundTrip(outputFile, null, null, variantRecords); } @Test public void testFormatModelVariantShreddingWithEncryption() throws IOException { - Schema variantSchema = VARIANT_SHREDDING_SCHEMA; - List variantRecords = variantShreddingRecords(variantSchema, 123456789, "string"); - ParquetFormatModel> model = - variantShreddingModel(variantSchema); + List variantRecords = variantShreddingRecords(123456789, "string"); + ParquetFormatModel> model = variantShreddingModel(); OutputFile encryptedFile = Files.localOutput(createTempFile(temp)); ByteBuffer fileDek = ByteBuffer.allocate(16); @@ -631,7 +496,7 @@ public void testFormatModelVariantShreddingWithEncryption() throws IOException { try (FileAppender appender = model .writeBuilder(EncryptedFiles.plainAsEncryptedOutput(encryptedFile)) - .schema(variantSchema) + .schema(VARIANT_SHREDDING_SCHEMA) .withFileEncryptionKey(fileDek) .withAADPrefix(aadPrefix) .setAll( @@ -647,56 +512,25 @@ public void testFormatModelVariantShreddingWithEncryption() throws IOException { assertThatThrownBy( () -> Parquet.read(encryptedFile.toInputFile()) - .project(variantSchema) + .project(VARIANT_SHREDDING_SCHEMA) .createReaderFunc( - fileSchema -> GenericParquetReaders.buildReader(variantSchema, fileSchema)) + fileSchema -> + GenericParquetReaders.buildReader(VARIANT_SHREDDING_SCHEMA, fileSchema)) .build() .iterator()) .isInstanceOf(ParquetCryptoRuntimeException.class) .hasMessage("Trying to read file with encrypted footer. No keys available"); - List writtenRecords; - try (CloseableIterable reader = - Parquet.read(encryptedFile.toInputFile()) - .project(variantSchema) - .withFileEncryptionKey(fileDek) - .withAADPrefix(aadPrefix) - .createReaderFunc( - fileSchema -> GenericParquetReaders.buildReader(variantSchema, fileSchema)) - .build()) { - writtenRecords = Lists.newArrayList(reader); - } - - assertThat(writtenRecords).hasSameSizeAs(variantRecords); - for (int i = 0; i < variantRecords.size(); i++) { - InternalTestHelpers.assertEquals( - variantSchema.asStruct(), variantRecords.get(i), writtenRecords.get(i)); - } + assertShreddedVariant( + encryptedFile, + FileDecryptionProperties.builder() + .withFooterKey(fileDek.array()) + .withAADPrefix(aadPrefix.array()) + .build(), + 123456789, + "string"); - try (ParquetReader rawReader = - ParquetReader.builder( - new GroupReadSupport(), new org.apache.hadoop.fs.Path(encryptedFile.location())) - .withDecryption( - FileDecryptionProperties.builder() - .withFooterKey(fileDek.array()) - .withAADPrefix(aadPrefix.array()) - .build()) - .build()) { - Group row = rawReader.read(); - Group variantData = row.getGroup("v", 0); - - assertThat(variantData.getFieldRepetitionCount("value")) - .as("value should be absent when fully shredded") - .isEqualTo(0); - - Group typedValue = variantData.getGroup("typed_value", 0); - assertThat(typedValue.getGroup("a", 0).getInteger("typed_value", 0)) - .as("typed_value.a should contain 123456789") - .isEqualTo(123456789); - assertThat(typedValue.getGroup("b", 0).getString("typed_value", 0)) - .as("typed_value.b should contain string") - .isEqualTo("string"); - } + assertRoundTrip(encryptedFile, fileDek, aadPrefix, variantRecords); } @ParameterizedTest @@ -722,6 +556,128 @@ public void testDefaultPathUsesCompressedSize() throws IOException { .isLessThan(trackedFile.splitOffsets().size()); } + private static ParquetFormatModel> variantModel( + VariantShreddingAnalyzer analyzer, + Function> transform) { + return ParquetFormatModel.create( + Record.class, + Void.class, + (icebergSchema, messageType, engineSchema) -> + GenericParquetWriter.create(icebergSchema, messageType), + (icebergSchema, fileSchema, engineSchema, idToConstant) -> + GenericParquetReaders.buildReader(icebergSchema, fileSchema), + analyzer, + transform); + } + + private static ParquetFormatModel> variantShreddingModel() { + VariantShreddingAnalyzer analyzer = + new VariantShreddingAnalyzer() { + @Override + protected List extractVariantValues(List rows, int idx) { + List values = Lists.newArrayList(); + for (Record row : rows) { + Object obj = row.get(idx); + if (obj instanceof Variant) { + values.add(((Variant) obj).value()); + } + } + return values; + } + + @Override + protected int resolveColumnIndex(Void engineSchema, String columnName) { + return VARIANT_SHREDDING_SCHEMA + .columns() + .indexOf(VARIANT_SHREDDING_SCHEMA.findField(columnName)); + } + }; + + return variantModel(analyzer, unused -> input -> input); + } + + private static List variantShreddingRecords(int aValue, String bValue) { + Variant variant = + VariantTestUtil.variant( + ImmutableMap.of("a", Variants.of(aValue), "b", Variants.of(bValue))); + GenericRecord record = GenericRecord.create(VARIANT_SHREDDING_SCHEMA); + return ImmutableList.of( + record.copy(ImmutableMap.of("id", 1L, "v", variant)), + record.copy(ImmutableMap.of("id", 2L, "v", variant)), + record.copy(ImmutableMap.of("id", 3L, "v", variant))); + } + + private static void assertShreddedVariant( + OutputFile file, FileDecryptionProperties decryption, int aValue, String bValue) + throws IOException { + ParquetReadOptions.Builder options = ParquetReadOptions.builder(); + if (decryption != null) { + options = options.withDecryption(decryption); + } + try (ParquetFileReader reader = + ParquetFileReader.open(ParquetIO.file(file.toInputFile()), options.build())) { + GroupType variantGroup = + reader.getFooter().getFileMetaData().getSchema().getType("v").asGroupType(); + assertThat(variantGroup.containsField("metadata")).isTrue(); + assertThat(variantGroup.containsField("value")).isTrue(); + assertThat(variantGroup.containsField("typed_value")).isTrue(); + + GroupType typedValueType = variantGroup.getType("typed_value").asGroupType(); + assertThat(typedValueType.containsField("a")).isTrue(); + assertThat(typedValueType.containsField("b")).isTrue(); + } + + ParquetReader.Builder builder = + ParquetReader.builder( + new GroupReadSupport(), new org.apache.hadoop.fs.Path(file.location())); + if (decryption != null) { + builder = builder.withDecryption(decryption); + } + try (ParquetReader rawReader = builder.build()) { + Group row; + while ((row = rawReader.read()) != null) { + Group variantData = row.getGroup("v", 0); + + assertThat(variantData.getFieldRepetitionCount("value")) + .as("value should be absent when fully shredded") + .isEqualTo(0); + + Group typedValue = variantData.getGroup("typed_value", 0); + assertThat(typedValue.getGroup("a", 0).getInteger("typed_value", 0)) + .as("typed_value.a should contain " + aValue) + .isEqualTo(aValue); + assertThat(typedValue.getGroup("b", 0).getString("typed_value", 0)) + .as("typed_value.b should contain " + bValue) + .isEqualTo(bValue); + } + } + } + + private static void assertRoundTrip( + OutputFile file, ByteBuffer fileDek, ByteBuffer aadPrefix, List expected) + throws IOException { + Parquet.ReadBuilder builder = + Parquet.read(file.toInputFile()) + .project(VARIANT_SHREDDING_SCHEMA) + .createReaderFunc( + fileSchema -> + GenericParquetReaders.buildReader(VARIANT_SHREDDING_SCHEMA, fileSchema)); + if (fileDek != null) { + builder = builder.withFileEncryptionKey(fileDek).withAADPrefix(aadPrefix); + } + + List writtenRecords; + try (CloseableIterable reader = builder.build()) { + writtenRecords = Lists.newArrayList(reader); + } + + assertThat(writtenRecords).hasSameSizeAs(expected); + for (int i = 0; i < expected.size(); i++) { + InternalTestHelpers.assertEquals( + VARIANT_SHREDDING_SCHEMA.asStruct(), expected.get(i), writtenRecords.get(i)); + } + } + // 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 {