ahmedabu98 commented on code in PR #40143:
URL: https://github.com/apache/beam/pull/40143#discussion_r4035622185
##########
sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/AddFilesTest.java:
##########
@@ -949,4 +963,671 @@ private DataWriter<Record> createWriter(String file,
@Nullable StructLike partit
private Record record(int id, String name, int age) {
return GenericRecord.create(icebergSchema).copy("id", id, "name", name,
"age", age);
}
+
+ // ---- ConvertToDataFile coverage check and pinned columns
+
+ private static final SchemaEvolutionConfig ADDITIONS =
+ SchemaEvolutionConfig.of(SchemaEvolutionOption.ALLOW_FIELD_ADDITION);
+
+ private PCollectionTuple convert(SchemaEvolutionConfig config, String...
files) {
+ PCollectionTuple out =
+ pipeline
+ .apply("Create Input", Create.of(Arrays.asList(files)))
+ .apply(
+ ParDo.of(
+ new AddFiles.ConvertToDataFile(
+ catalogConfig, tableId.toString(), null, null,
null, null, config))
+ .withOutputTags(
+ AddFiles.ConvertToDataFile.DATA_FILES,
+ TupleTagList.of(AddFiles.ConvertToDataFile.ERRORS)));
+
out.get(AddFiles.ConvertToDataFile.ERRORS).setRowSchema(AddFiles.ERROR_SCHEMA);
+ return out;
+ }
+
+ private String writeWithSchema(String name, Schema schema, Record...
records) throws IOException {
+ String file = root + name;
+ DataWriter<Record> writer =
+ Parquet.writeData(Files.localOutput(file))
+ .schema(schema)
+ .withSpec(PartitionSpec.unpartitioned())
+ .createWriterFunc(GenericParquetWriter::create)
+ .build();
+ try {
+ for (Record record : records) {
+ writer.write(record);
+ }
+ } finally {
+ writer.close();
+ }
+ return file;
+ }
+
+ private String writeOneRecord(String name) throws IOException {
+ String file = root + name;
+ DataWriter<Record> writer = createWriter(file);
+ writer.write(record(1, "a", 1));
+ writer.close();
+ return file;
+ }
+
+ private static final Schema WIDER =
+ new Schema(
+ Types.NestedField.required(1, "id", Types.IntegerType.get()),
+ Types.NestedField.required(2, "name", Types.StringType.get()),
+ Types.NestedField.required(3, "age", Types.IntegerType.get()),
+ Types.NestedField.optional(4, "email", Types.StringType.get()));
+
+ private String writeWider(String name) throws IOException {
+ Record record = GenericRecord.create(WIDER);
+ record.setField("id", 1);
+ record.setField("name", "a");
+ record.setField("age", 1);
+ record.setField("email", "e");
+ return writeWithSchema(name, WIDER, record);
+ }
+
+ private void assertSingleError(PCollectionTuple out, String file, String
contains) {
+ PAssert.that(out.get(AddFiles.ConvertToDataFile.DATA_FILES)).empty();
+ PAssert.that(out.get(AddFiles.ConvertToDataFile.ERRORS))
+ .satisfies(
+ rows -> {
+ Row row = Iterables.getOnlyElement(rows);
+ assertEquals(file, row.getString("file"));
+ assertThat(row.getString("error"), containsString(contains));
+ return null;
+ });
+ }
+
+ private void assertRegisters(PCollectionTuple out, long files) {
+ PAssert.that(out.get(AddFiles.ConvertToDataFile.ERRORS)).empty();
+
PAssert.thatSingleton(out.get(AddFiles.ConvertToDataFile.DATA_FILES).apply(Count.globally()))
+ .isEqualTo(files);
+ }
+
+ @Test
+ public void testEmptySchemaTableWarnsAndStillRegisters() throws Exception {
+ catalog.createTable(tableId, new Schema());
+ String file = writeOneRecord("data.parquet");
+
+ PCollectionTuple out = convert(SchemaEvolutionConfig.disabled(), file);
+
+ assertRegisters(out, 1);
+ pipeline.run().waitUntilFinish();
+ logs.verifyWarn("has no columns");
+ }
+
+ /** Files without embedded field ids against a zero-column table still
register. */
+ @Test
+ public void testEmptySchemaTableWithIdLessParquetStillRegisters() throws
Exception {
+ catalog.createTable(tableId, new Schema());
+ File file = new File(temp.getRoot(), "idless.parquet");
+ org.apache.avro.Schema avro =
+ org.apache.avro.SchemaBuilder.record("r")
+ .fields()
+ .requiredInt("id")
+ .optionalString("name")
+ .name("address")
+ .type()
+ .record("address")
+ .fields()
+ .optionalString("city")
+ .endRecord()
+ .noDefault()
+ .endRecord();
+ try (org.apache.parquet.hadoop.ParquetWriter<Object> writer =
+ org.apache.parquet.avro.AvroParquetWriter.builder(
+ new org.apache.hadoop.fs.Path(file.getAbsolutePath()))
+ .withSchema(avro)
+ .build()) {
+ org.apache.avro.generic.GenericData.Record record =
+ new org.apache.avro.generic.GenericData.Record(avro);
+ record.put("id", 1);
+ record.put("name", "a");
+ org.apache.avro.generic.GenericData.Record address =
+ new
org.apache.avro.generic.GenericData.Record(avro.getField("address").schema());
+ address.put("city", "c");
+ record.put("address", address);
+ writer.write(record);
+ }
+
+ PCollectionTuple out = convert(SchemaEvolutionConfig.disabled(),
file.getAbsolutePath());
+
+ assertRegisters(out, 1);
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testCoveredFileRegistersWithEvolutionEnabled() throws Exception {
+ catalog.createTable(tableId, icebergSchema);
+ String file = writeOneRecord("data.parquet");
+
+ PCollectionTuple out = convert(ADDITIONS, file);
+
+ assertRegisters(out, 1);
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testUncoveredFileRoutesToErrorsWhenEvolutionEnabled() throws
Exception {
+ catalog.createTable(tableId, icebergSchema);
+ String file = writeWider("wider.parquet");
+
+ PCollectionTuple out = convert(ADDITIONS, file);
+
+ assertSingleError(out, file, "does not cover the file");
+ PAssert.that(out.get(AddFiles.ConvertToDataFile.ERRORS))
+ .satisfies(
+ rows -> {
+ assertThat(
+ Iterables.getOnlyElement(rows).getString("error"),
+ containsString("add optional email string"));
+ return null;
+ });
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testExtraColumnsRegisterWhenEvolutionDisabled() throws Exception
{
+ catalog.createTable(tableId, icebergSchema);
+ String file = writeWider("wider.parquet");
+
+ PCollectionTuple out = convert(SchemaEvolutionConfig.disabled(), file);
+
+ assertRegisters(out, 1);
+ pipeline.run().waitUntilFinish();
+ }
+
+ @Test
+ public void testUnreadableSchemaRoutesToErrorsWithConverterMessage() throws
Exception {
+ catalog.createTable(tableId, icebergSchema);
+ // a legacy unannotated repeated field: readable Parquet, rejected by
Iceberg's converter
+ org.apache.parquet.schema.MessageType legacy =
+ org.apache.parquet.schema.Types.buildMessage()
+
.required(org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32)
+ .named("id")
+
.repeated(org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName.INT32)
+ .named("vals")
+ .named("root");
+ File legacyFile = new File(temp.getRoot(), "legacy.parquet");
+ try
(org.apache.parquet.hadoop.ParquetWriter<org.apache.parquet.example.data.Group>
writer =
+ org.apache.parquet.hadoop.example.ExampleParquetWriter.builder(
+ new org.apache.hadoop.fs.Path(legacyFile.getAbsolutePath()))
+ .withType(legacy)
+ .build()) {
+ org.apache.parquet.example.data.Group group =
+ new
org.apache.parquet.example.data.simple.SimpleGroupFactory(legacy).newGroup();
+ group.add("id", 1);
+ group.add("vals", 2);
+ writer.write(group);
+ }
+ String file = legacyFile.getAbsolutePath();
+
+ PCollectionTuple out = convert(ADDITIONS, file);
+
+ assertSingleError(out, file,
AddFiles.ConvertToDataFile.UNREADABLE_SCHEMA_ERROR);
+ PAssert.that(out.get(AddFiles.ConvertToDataFile.ERRORS))
+ .satisfies(
+ rows -> {
+ assertThat(
+ Iterables.getOnlyElement(rows).getString("error"),
+ containsString("repetition REPEATED"));
+ return null;
+ });
+ pipeline.run().waitUntilFinish();
+ }
+
+ // ---- pinned columns
+ //
+ // A pin is enforced in two layers. A pinned column the table holds as
REQUIRED is protected by
+ // the coverage check: a file that declares it optional with nulls, or lacks
it, needs a
+ // relaxation of a pinned column, which SchemaDelta refuses, so the file is
routed there
+ // (testPinnedRequiredColumnIsProtectedByCoverage). The per-file pin walk is
reached only for
+ // pinned columns the table holds as OPTIONAL: columns the pre-pass added
(it never creates them
+ // required) or pre-existing optional ones. The tests below therefore pin
optional columns.
Review Comment:
> it never creates them required
From what I'm seeing, this is only true if the table already exists and it
needs to add new pinned column. But if it needs to create the table itself then
it will create those pinned columns required right? If that's intended behavior
maybe we should mention it in javadoc somewhere
--
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]
