diff --git a/CHANGES.md b/CHANGES.md index d1910513da58..5b0fbed21141 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -97,6 +97,7 @@ * (Go) Wrap last dataWriter/timerWriter send errors on Close ([#40405](https://github.com/apache/beam/issues/40405)). * (Go) Fixed a deadlock recreating a data channel while holding the channel lock ([#40414](https://github.com/apache/beam/issues/40414)). * (Java) Fixed the declared schema of the error output of the Kafka write SchemaTransform, which wrapped the error schema a second time and did not match the rows it emits ([#39760](https://github.com/apache/beam/issues/39760)). +* (Java) ParquetIO now uses the schema passed to `read()`/`readFiles()` as the Avro reader schema and matches file columns to its fields by name. Reading files that predate a newly added field no longer fails with `ArrayIndexOutOfBoundsException`, and a schema that omits some of the file's columns no longer returns values from the wrong column ([#27234](https://github.com/apache/beam/issues/27234)). * (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed in recent releases, which broke builds of Go modules depending on a current `genproto` version ([#40018](https://github.com/apache/beam/issues/40018)). * (Java) BigQueryIO now treats a 404 when deleting a temporary table or dataset as success, so a replayed work item whose earlier attempt already deleted it no longer retries forever ([#24997](https://github.com/apache/beam/issues/24997)). * (Java) IcebergIO now writes rows containing `EnumerationType` (proto enum) fields as strings, instead of throwing `Unsupported Beam logical type Enum` ([#40299](https://github.com/apache/beam/issues/40299)). diff --git a/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java b/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java index 356cd3514ac1..75fbab0ff76f 100644 --- a/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java +++ b/sdks/java/io/parquet/src/main/java/org/apache/beam/sdk/io/parquet/ParquetIO.java @@ -29,9 +29,11 @@ import java.nio.channels.SeekableByteChannel; import java.nio.channels.WritableByteChannel; import java.util.ArrayList; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Set; import javax.annotation.Nullable; import org.apache.avro.Schema; import org.apache.avro.generic.GenericData; @@ -91,6 +93,7 @@ import org.apache.parquet.io.SeekableInputStream; import org.apache.parquet.io.api.RecordMaterializer; import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.Type; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -114,6 +117,12 @@ * *

As {@link Read} is based on {@link FileIO}, it supports any filesystem (hdfs, ...). * + *

ParquetIO uses the schema passed to {@link #read(Schema)} or {@link #readFiles(Schema)} as the + * Avro reader schema and matches file columns to its fields by name or alias. This lets you read + * files written with an older version of the schema: a field the file lacks gets its default value, + * and a top-level column the schema lacks is not read. With {@link Read#withProjection(Schema, + * Schema)}, the encoder schema takes this role. + * *

When using schemas created via reflection, it may be useful to generate {@link GenericRecord} * instances rather than instances of the class associated with the schema. {@link Read} and {@link * ReadFiles} provide {@link ParquetIO.Read#withAvroDataModel(GenericData)} allowing implementations @@ -554,7 +563,7 @@ public PCollection expand(PCollection input) { checkArgument(!isGenericRecordOutput(), "Parse can't be used for reading as GenericRecord."); return input - .apply(ParDo.of(new SplitReadFn<>(null, null, getParseFn(), getConfiguration()))) + .apply(ParDo.of(new SplitReadFn<>(null, null, null, getParseFn(), getConfiguration()))) .setCoder(inferCoder(input.getPipeline().getCoderRegistry())); } @@ -697,11 +706,20 @@ public PCollection expand(PCollection input) { new SplitReadFn<>( getAvroDataModel(), getProjectionSchema(), + getOutputSchema(), GenericRecordPassthroughFn.create(), getConfiguration()))) .setCoder(getCollectionCoder()); } + /** + * Returns the schema of the output records. It is used both as the Avro reader schema, so that + * file records are resolved against it, and as the schema of the output coder. + */ + private Schema getOutputSchema() { + return getProjectionSchema() != null ? getEncoderSchema() : getSchema(); + } + @Override public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); @@ -727,7 +745,7 @@ public void populateDisplayData(DisplayData.Builder builder) { * AvroCoder} when not using Beam schema. */ private Coder getCollectionCoder() { - Schema coderSchema = getProjectionSchema() != null ? getEncoderSchema() : getSchema(); + Schema coderSchema = getOutputSchema(); return getInferBeamSchema() ? AvroUtils.schemaCoder(coderSchema) : AvroCoder.of(coderSchema); } @@ -736,8 +754,11 @@ private Coder getCollectionCoder() { static final class SplitReadFn extends DoFn { private final Class modelClass; private final String requestSchemaString; + private final String readSchemaString; // Default initial splitting the file into blocks of 64MB. Unit of SPLIT_LIMIT is byte. private static final long SPLIT_LIMIT = 64000000; + // Same key as the private AvroReadSupport.AVRO_READ_SCHEMA. + private static final String AVRO_READ_SCHEMA = "parquet.avro.read.schema"; private @Nullable final SerializableConfiguration configuration; @@ -746,11 +767,13 @@ static final class SplitReadFn extends DoFn { SplitReadFn( GenericData model, Schema requestSchema, + Schema readSchema, SerializableFunction parseFn, @Nullable SerializableConfiguration configuration) { this.modelClass = model != null ? model.getClass() : null; this.requestSchemaString = requestSchema != null ? requestSchema.toString() : null; + this.readSchemaString = readSchema != null ? readSchema.toString() : null; this.parseFn = checkNotNull(parseFn, "GenericRecord parse function can't be null"); this.configuration = configuration; } @@ -780,6 +803,11 @@ public void processElement( AvroReadSupport.setRequestedProjection( conf, new Schema.Parser().parse(requestSchemaString)); } + // Resolve file records against the output schema, unless the user already set a read schema + // through the configuration. + if (readSchemaString != null && conf.get(AVRO_READ_SCHEMA) == null) { + AvroReadSupport.setAvroReadSchema(conf, new Schema.Parser().parse(readSchemaString)); + } ParquetReadOptions options = HadoopReadOptions.builder(conf).build(); try (ParquetFileReader reader = ParquetFileReader.open(new BeamParquetInputFile(file.openSeekable()), options)) { @@ -794,6 +822,14 @@ public void processElement( hadoopConf, Maps.transformValues(fileMetadata, ImmutableSet::of), fileSchema)); + String avroReadSchema = hadoopConf.get(AVRO_READ_SCHEMA); + if (avroReadSchema != null + && hadoopConf.get(AvroReadSupport.AVRO_REQUESTED_PROJECTION) == null) { + readContext = + new ReadSupport.ReadContext( + pruneToReadSchema(fileSchema, new Schema.Parser().parse(avroReadSchema)), + readContext.getReadSupportMetadata()); + } ColumnIOFactory columnIOFactory = new ColumnIOFactory(parquetFileMetadata.getCreatedBy()); RecordMaterializer recordConverter = @@ -877,6 +913,31 @@ public Configuration getConfWithModelClass() throws ReflectiveOperationException return conf; } + /** + * Drops the top-level columns of {@code fileSchema} that have no matching field (by name or + * alias) in {@code readSchema}. The Avro record converter fails on columns missing from the + * read schema, so without this a read schema with fewer fields than the file could not be + * used. Returns {@code fileSchema} unchanged if no columns would remain. + */ + @VisibleForTesting + static MessageType pruneToReadSchema(MessageType fileSchema, Schema readSchema) { + Set names = new HashSet<>(); + for (Schema.Field field : readSchema.getFields()) { + names.add(field.name()); + names.addAll(field.aliases()); + } + List fields = new ArrayList<>(); + for (Type field : fileSchema.getFields()) { + if (names.contains(field.getName())) { + fields.add(field); + } + } + if (fields.isEmpty() || fields.size() == fileSchema.getFieldCount()) { + return fileSchema; + } + return new MessageType(fileSchema.getName(), fields); + } + @GetInitialRestriction public OffsetRange getInitialRestriction(@Element ReadableFile file) throws Exception { try (ParquetFileReader reader = getParquetFileReader(file)) { diff --git a/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java b/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java index 97f96452760f..a7339c360c96 100644 --- a/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java +++ b/sdks/java/io/parquet/src/test/java/org/apache/beam/sdk/io/parquet/ParquetIOTest.java @@ -56,6 +56,7 @@ import org.apache.beam.sdk.transforms.display.DisplayData; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.hadoop.conf.Configuration; import org.apache.parquet.filter2.predicate.FilterApi; import org.apache.parquet.filter2.predicate.FilterPredicate; @@ -65,6 +66,8 @@ import org.apache.parquet.hadoop.metadata.ParquetMetadata; import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.parquet.io.api.Binary; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.MessageTypeParser; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TemporaryFolder; @@ -158,7 +161,7 @@ public void testBlockTracker() { public void testSplitBlockWithLimit() { ParquetIO.ReadFiles.SplitReadFn testFn = new ParquetIO.ReadFiles.SplitReadFn<>( - null, null, ParquetIO.GenericRecordPassthroughFn.create(), null); + null, null, null, ParquetIO.GenericRecordPassthroughFn.create(), null); ArrayList blockList = new ArrayList<>(); ArrayList rangeList; BlockMetaData testBlock = mock(BlockMetaData.class); @@ -495,6 +498,216 @@ public void testWriteAndReadWithConfiguration() { readPipeline.run().waitUntilFinish(); } + private static final Schema SCHEMA_WITH_NULLABLE_FIELD = + new Schema.Parser() + .parse( + "{" + + "\"type\":\"record\", " + + "\"name\":\"testrecord\"," + + "\"fields\":[" + + " {\"name\":\"name\",\"type\":\"string\"}," + + " {\"name\":\"id\",\"type\":\"string\"}," + + " {\"name\":\"country\",\"type\":[\"null\",\"string\"],\"default\":null}" + + " ]" + + "}"); + + private static final Schema SCHEMA_WITH_DEFAULT_FIELD = + new Schema.Parser() + .parse( + "{" + + "\"type\":\"record\", " + + "\"name\":\"testrecord\"," + + "\"fields\":[" + + " {\"name\":\"name\",\"type\":\"string\"}," + + " {\"name\":\"id\",\"type\":\"string\"}," + + " {\"name\":\"country\",\"type\":\"string\",\"default\":\"unknown\"}" + + " ]" + + "}"); + + private void writeGenericRecords(List records) { + mainPipeline + .apply(Create.of(records).withCoder(AvroCoder.of(SCHEMA))) + .apply( + FileIO.write() + .via(ParquetIO.sink(SCHEMA)) + .to(temporaryFolder.getRoot().getAbsolutePath())); + mainPipeline.run().waitUntilFinish(); + } + + private static List withCountry( + List records, Schema schema, String country) { + return records.stream() + .map( + record -> + (GenericRecord) + new GenericRecordBuilder(schema) + .set("name", record.get("name")) + .set("id", record.get("id")) + .set("country", country) + .build()) + .collect(toList()); + } + + @Test + public void testReadWithAddedNullableField() { + List records = generateGenericRecords(1000); + writeGenericRecords(records); + + PCollection readBack = + readPipeline.apply( + ParquetIO.read(SCHEMA_WITH_NULLABLE_FIELD) + .from(temporaryFolder.getRoot().getAbsolutePath() + "/*")); + PAssert.that(readBack) + .containsInAnyOrder(withCountry(records, SCHEMA_WITH_NULLABLE_FIELD, null)); + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testReadFilesWithAddedFieldWithDefault() { + List records = generateGenericRecords(1000); + writeGenericRecords(records); + + PCollection readBack = + readPipeline + .apply(FileIO.match().filepattern(temporaryFolder.getRoot().getAbsolutePath() + "/*")) + .apply(FileIO.readMatches()) + .apply(ParquetIO.readFiles(SCHEMA_WITH_DEFAULT_FIELD)); + PAssert.that(readBack) + .containsInAnyOrder(withCountry(records, SCHEMA_WITH_DEFAULT_FIELD, "unknown")); + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testReadWithSubsetSchemaWithoutProjection() { + // REQUESTED_SCHEMA drops the first column, so matching fields by position would return names + // as ids. + List records = generateGenericRecords(1000); + writeGenericRecords(records); + + List expected = + records.stream() + .map( + record -> + (GenericRecord) + new GenericRecordBuilder(REQUESTED_SCHEMA) + .set("id", record.get("id")) + .build()) + .collect(toList()); + PCollection readBack = + readPipeline.apply( + ParquetIO.read(REQUESTED_SCHEMA) + .from(temporaryFolder.getRoot().getAbsolutePath() + "/*")); + PAssert.that(readBack).containsInAnyOrder(expected); + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testReadWithProjectionAndAddedField() { + Schema encoderSchema = + new Schema.Parser() + .parse( + "{" + + "\"type\":\"record\", " + + "\"name\":\"testrecord\"," + + "\"fields\":[" + + " {\"name\":\"name\",\"type\":[\"string\",\"null\"]}," + + " {\"name\":\"id\",\"type\":\"string\"}," + + " {\"name\":\"country\",\"type\":[\"null\",\"string\"],\"default\":null}" + + " ]" + + "}"); + writeGenericRecords(generateGenericRecords(1000)); + + List expected = new ArrayList<>(); + for (int i = 0; i < 1000; i++) { + expected.add( + new GenericRecordBuilder(encoderSchema) + .set("name", null) + .set("id", Integer.toString(i)) + .set("country", null) + .build()); + } + PCollection readBack = + readPipeline.apply( + ParquetIO.read(SCHEMA) + .from(temporaryFolder.getRoot().getAbsolutePath() + "/*") + .withProjection(REQUESTED_SCHEMA, encoderSchema)); + PAssert.that(readBack).containsInAnyOrder(expected); + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testReadSchemaFromConfigurationIsNotOverridden() { + Schema configuredReadSchema = + new Schema.Parser() + .parse( + "{" + + "\"type\":\"record\", " + + "\"name\":\"testrecord\"," + + "\"fields\":[" + + " {\"name\":\"name\",\"type\":\"string\"}," + + " {\"name\":\"id\",\"type\":\"string\"}," + + " {\"name\":\"country\",\"type\":[\"string\",\"null\"]," + + "\"default\":\"configured\"}" + + " ]" + + "}"); + List records = generateGenericRecords(1000); + writeGenericRecords(records); + + PCollection readBack = + readPipeline.apply( + ParquetIO.read(SCHEMA_WITH_NULLABLE_FIELD) + .from(temporaryFolder.getRoot().getAbsolutePath() + "/*") + .withConfiguration( + ImmutableMap.of("parquet.avro.read.schema", configuredReadSchema.toString()))); + PAssert.that(readBack) + .containsInAnyOrder(withCountry(records, SCHEMA_WITH_NULLABLE_FIELD, "configured")); + readPipeline.run().waitUntilFinish(); + } + + @Test + public void testPruneToReadSchema() { + MessageType fileSchema = + MessageTypeParser.parseMessageType( + "message testrecord {" + + " required binary name (STRING);" + + " required binary id (STRING);" + + " optional binary country (STRING);" + + "}"); + Schema readSchema = + new Schema.Parser() + .parse( + "{" + + "\"type\":\"record\", " + + "\"name\":\"testrecord\"," + + "\"fields\":[" + + " {\"name\":\"id\",\"type\":\"string\"}," + + " {\"name\":\"fullName\",\"aliases\":[\"name\"],\"type\":\"string\"}," + + " {\"name\":\"age\",\"type\":[\"null\",\"int\"],\"default\":null}" + + " ]" + + "}"); + + MessageType pruned = ParquetIO.ReadFiles.SplitReadFn.pruneToReadSchema(fileSchema, readSchema); + + assertEquals( + MessageTypeParser.parseMessageType( + "message testrecord {" + + " required binary name (STRING);" + + " required binary id (STRING);" + + "}"), + pruned); + assertEquals( + fileSchema, + ParquetIO.ReadFiles.SplitReadFn.pruneToReadSchema(fileSchema, SCHEMA_WITH_NULLABLE_FIELD)); + assertEquals( + fileSchema, + ParquetIO.ReadFiles.SplitReadFn.pruneToReadSchema( + fileSchema, + new Schema.Parser() + .parse( + "{\"type\":\"record\",\"name\":\"other\",\"fields\":" + + "[{\"name\":\"unrelated\",\"type\":\"string\"}]}"))); + } + @Test public void testWriteAndReadFilesAsJsonForUnknownSchemaWithConfiguration() { List records = generateGenericRecords(10);