wombatu-kun commented on code in PR #17469:
URL: https://github.com/apache/iceberg/pull/17469#discussion_r3698272103
##########
parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java:
##########
@@ -480,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<Record, Void, ParquetValueReader<?>>
variantShreddingModel(
Review Comment:
variantShreddingModel and variantShreddingRecords take a Schema that is only
ever VARIANT_SHREDDING_SCHEMA, and each test adds a variantSchema alias just to
pass it. Drop the parameter and read the constant inside the helpers.
##########
parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java:
##########
@@ -503,40 +507,43 @@ protected List<VariantValue>
extractVariantValues(List<Record> 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<Void, UnaryOperator<Record>>) unused -> input -> input);
+ }
+
+ private static List<Record> 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)));
Review Comment:
The new variantShreddingRecords helper hand-builds the metadata buffer,
object buffer and Variant, which VariantTestUtil.variant(Map) already does in
one call. Replace those four statements with
VariantTestUtil.variant(ImmutableMap.of("a", Variants.of(aValue), "b",
Variants.of(bValue))).
##########
parquet/src/test/java/org/apache/iceberg/parquet/TestParquetDataWriter.java:
##########
@@ -607,6 +614,91 @@ protected int resolveColumnIndex(Void engineSchema, String
columnName) {
}
}
+ @Test
+ public void testFormatModelVariantShreddingWithEncryption() throws
IOException {
+ Schema variantSchema = VARIANT_SHREDDING_SCHEMA;
+ List<Record> variantRecords = variantShreddingRecords(variantSchema,
123456789, "string");
+ ParquetFormatModel<Record, Void, ParquetValueReader<?>> model =
+ variantShreddingModel(variantSchema);
+
+ 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<Record> 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<Record> writtenRecords;
+ try (CloseableIterable<Record> 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 (ParquetReader<Group> rawReader =
Review Comment:
The raw-group block here is a line-for-line copy of the one in
testFormatModelVariantShreddingRoundTrip apart from the withDecryption call and
the two literals. Fold it into a helper next to variantShreddingModel and
variantShreddingRecords that takes the file, the decryption properties, and the
expected a/b values.
--
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]
