diff --git a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroDeserializationSchema.java b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroDeserializationSchema.java index 226f47a339b213..3d61d709b46f4c 100644 --- a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroDeserializationSchema.java +++ b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroDeserializationSchema.java @@ -71,7 +71,22 @@ public static AvroDeserializationSchema forGeneric(Schema schema) */ public static AvroDeserializationSchema forGeneric( Schema schema, AvroEncoding encoding) { - return new AvroDeserializationSchema<>(GenericRecord.class, schema, encoding); + return forGeneric(schema, null, encoding); + } + + /** + * Creates {@link AvroDeserializationSchema} that produces {@link GenericRecord} using the + * provided reader schema and resolves the input against the given writer schema. + * + * @param reader reader schema of produced records + * @param writer writer schema the input was serialized with, or {@code null} to reuse the + * reader schema + * @param encoding Avro serialization approach to use for decoding + * @return deserialized record in form of {@link GenericRecord} + */ + static AvroDeserializationSchema forGeneric( + Schema reader, @Nullable Schema writer, AvroEncoding encoding) { + return new AvroDeserializationSchema<>(GenericRecord.class, reader, writer, encoding); } /** @@ -107,6 +122,9 @@ public static AvroDeserializationSchema forSpecifi /** Schema in case of GenericRecord for serialization purpose. */ private final String schemaString; + /** Writer schema the input was serialized with, or {@code null} to reuse the reader schema. */ + private final String writerSchemaString; + /** Reader that deserializes byte array into a record. */ private transient GenericDatumReader datumReader; @@ -122,6 +140,9 @@ public static AvroDeserializationSchema forSpecifi /** Avro schema for the reader. */ private transient Schema reader; + /** Avro schema for the writer. */ + private transient Schema writer; + /** * Creates a Avro deserialization schema. * @@ -133,14 +154,31 @@ public static AvroDeserializationSchema forSpecifi */ AvroDeserializationSchema( Class recordClazz, @Nullable Schema reader, AvroEncoding encoding) { + this(recordClazz, reader, null, encoding); + } + + /** + * Creates a Avro deserialization schema. + * + * @param recordClazz class to which deserialize. Should be one of: {@link + * org.apache.avro.specific.SpecificRecord}, {@link org.apache.avro.generic.GenericRecord}. + * @param reader reader's Avro schema. Should be provided if recordClazz is {@link + * GenericRecord} + * @param writer writer's Avro schema the input was serialized with. Used to resolve the input + * against the reader schema, e.g. to read a subset of the written fields. + * @param encoding encoding approach to use. Identifies the Avro decoder class to use. + */ + AvroDeserializationSchema( + Class recordClazz, + @Nullable Schema reader, + @Nullable Schema writer, + AvroEncoding encoding) { Preconditions.checkNotNull(recordClazz, "Avro record class must not be null."); this.recordClazz = recordClazz; this.reader = reader; - if (reader != null) { - this.schemaString = reader.toString(); - } else { - this.schemaString = null; - } + this.schemaString = reader != null ? reader.toString() : null; + this.writer = writer; + this.writerSchemaString = writer != null ? writer.toString() : null; this.encoding = encoding; } @@ -152,6 +190,11 @@ Schema getReaderSchema() { return reader; } + @Nullable + Schema getWriterSchema() { + return writer; + } + MutableByteArrayInputStream getInputStream() { return inputStream; } @@ -173,9 +216,15 @@ public T deserialize(@Nullable byte[] message) throws IOException { checkAvroInitialized(); inputStream.setBuffer(message); Schema readerSchema = getReaderSchema(); + Schema writerSchema = getWriterSchema(); GenericDatumReader datumReader = getDatumReader(); - datumReader.setSchema(readerSchema); + if (writerSchema != null) { + datumReader.setSchema(writerSchema); + datumReader.setExpected(readerSchema); + } else { + datumReader.setSchema(readerSchema); + } if (encoding == AvroEncoding.JSON) { ((JsonDecoder) this.decoder).configure(inputStream); @@ -199,6 +248,9 @@ void checkAvroInitialized() throws IOException { this.reader = AvroFactory.extractAvroSpecificSchema(recordClazz, specificData); } else { this.reader = new Schema.Parser().parse(schemaString); + if (writerSchemaString != null) { + this.writer = new Schema.Parser().parse(writerSchemaString); + } GenericData genericData = new GenericData(cl); this.datumReader = new GenericDatumReader<>(null, this.reader, genericData); } @@ -206,7 +258,9 @@ void checkAvroInitialized() throws IOException { this.inputStream = new MutableByteArrayInputStream(); if (encoding == AvroEncoding.JSON) { - this.decoder = DecoderFactory.get().jsonDecoder(getReaderSchema(), inputStream); + final Schema writerSchema = getWriterSchema(); + final Schema decoderSchema = writerSchema == null ? getReaderSchema() : writerSchema; + this.decoder = DecoderFactory.get().jsonDecoder(decoderSchema, inputStream); } else { this.decoder = DecoderFactory.get().binaryDecoder(inputStream, null); } @@ -236,11 +290,13 @@ public boolean equals(Object o) { return false; } AvroDeserializationSchema that = (AvroDeserializationSchema) o; - return recordClazz.equals(that.recordClazz) && Objects.equals(reader, that.reader); + return recordClazz.equals(that.recordClazz) + && Objects.equals(reader, that.reader) + && Objects.equals(writer, that.writer); } @Override public int hashCode() { - return Objects.hash(recordClazz, reader); + return Objects.hash(recordClazz, reader, writer); } } diff --git a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFormatFactory.java b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFormatFactory.java index eda1cc1a898ff8..d6b71ca7d4ed00 100644 --- a/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFormatFactory.java +++ b/flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFormatFactory.java @@ -25,6 +25,7 @@ import org.apache.flink.configuration.ConfigOption; import org.apache.flink.configuration.ReadableConfig; import org.apache.flink.formats.avro.AvroFormatOptions.AvroEncoding; +import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter; import org.apache.flink.table.connector.ChangelogMode; import org.apache.flink.table.connector.Projection; import org.apache.flink.table.connector.format.DecodingFormat; @@ -72,11 +73,22 @@ public DeserializationSchema createRuntimeDecoder( int[][] projections) { final DataType producedDataType = Projection.of(projections).project(physicalDataType); - final RowType rowType = (RowType) producedDataType.getLogicalType(); + final RowType producedRowType = (RowType) producedDataType.getLogicalType(); + final RowType physicalRowType = (RowType) physicalDataType.getLogicalType(); final TypeInformation rowDataTypeInfo = context.createTypeInformation(producedDataType); return new AvroRowDataDeserializationSchema( - rowType, rowDataTypeInfo, encoding, legacyTimestampMapping); + AvroDeserializationSchema.forGeneric( + AvroSchemaConverter.convertToSchema( + producedRowType, legacyTimestampMapping), + producedRowType.equals(physicalRowType) + ? null + : AvroSchemaConverter.convertToSchema( + physicalRowType, legacyTimestampMapping), + encoding), + AvroToRowDataConverters.createRowConverter( + producedRowType, legacyTimestampMapping), + rowDataTypeInfo); } @Override diff --git a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFormatFactoryTest.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFormatFactoryTest.java index 0b9d572a13b07a..4a8d9da7118856 100644 --- a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFormatFactoryTest.java +++ b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroFormatFactoryTest.java @@ -21,9 +21,11 @@ import org.apache.flink.api.common.serialization.DeserializationSchema; import org.apache.flink.api.common.serialization.SerializationSchema; import org.apache.flink.formats.avro.AvroFormatOptions.AvroEncoding; +import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter; import org.apache.flink.table.api.DataTypes; import org.apache.flink.table.catalog.Column; import org.apache.flink.table.catalog.ResolvedSchema; +import org.apache.flink.table.connector.format.ProjectableDecodingFormat; import org.apache.flink.table.connector.sink.DynamicTableSink; import org.apache.flink.table.connector.source.DynamicTableSource; import org.apache.flink.table.data.RowData; @@ -31,13 +33,21 @@ import org.apache.flink.table.factories.utils.FactoryMocks; import org.apache.flink.table.runtime.connector.source.ScanRuntimeProviderContext; import org.apache.flink.table.runtime.typeutils.InternalTypeInfo; +import org.apache.flink.table.types.DataType; import org.apache.flink.table.types.logical.RowType; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.avro.io.Encoder; +import org.apache.avro.io.EncoderFactory; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.EnumSource; import org.junit.jupiter.params.provider.ValueSource; +import java.io.ByteArrayOutputStream; import java.util.HashMap; import java.util.Map; @@ -54,6 +64,13 @@ class AvroFormatFactoryTest { Column.physical("c", DataTypes.BOOLEAN()), Column.physical("d", DataTypes.TIMESTAMP(3))); + private static final ResolvedSchema PROJECTION_SCHEMA = + ResolvedSchema.of( + Column.physical("user_id", DataTypes.BIGINT().notNull()), + Column.physical("name", DataTypes.STRING().notNull()), + Column.physical("event_id", DataTypes.BIGINT().notNull()), + Column.physical("payload", DataTypes.STRING().notNull())); + private static final ResolvedSchema NEW_SCHEMA = ResolvedSchema.of( Column.physical("a", DataTypes.STRING()), @@ -77,7 +94,7 @@ void testSeDeSchema(AvroEncoding encoding) { new AvroRowDataDeserializationSchema( ROW_TYPE, InternalTypeInfo.of(ROW_TYPE), encoding); - final Map options = getAllOptions(true); + final Map options = getAllOptions(true, encoding); final DynamicTableSource actualSource = FactoryMocks.createTableSource(SCHEMA, options); assertThat(actualSource).isInstanceOf(TestDynamicTableFactory.DynamicTableSourceMock.class); @@ -104,6 +121,51 @@ void testSeDeSchema(AvroEncoding encoding) { assertThat(actualSer).isEqualTo(expectedSer); } + @ParameterizedTest + @EnumSource(AvroEncoding.class) + void testProjectionPushDown(AvroEncoding encoding) throws Exception { + final Map options = getAllOptions(true, encoding); + final DynamicTableSource actualSource = + FactoryMocks.createTableSource(PROJECTION_SCHEMA, options); + assertThat(actualSource).isInstanceOf(TestDynamicTableFactory.DynamicTableSourceMock.class); + final TestDynamicTableFactory.DynamicTableSourceMock scanSourceMock = + (TestDynamicTableFactory.DynamicTableSourceMock) actualSource; + assertThat(scanSourceMock.valueFormat).isInstanceOf(ProjectableDecodingFormat.class); + + final ProjectableDecodingFormat> projectableFormat = + (ProjectableDecodingFormat>) + scanSourceMock.valueFormat; + final DataType physicalDataType = PROJECTION_SCHEMA.toPhysicalRowDataType(); + final DeserializationSchema deserializationSchema = + projectableFormat.createRuntimeDecoder( + ScanRuntimeProviderContext.INSTANCE, + physicalDataType, + new int[][] {{1}, {2}}); + deserializationSchema.open(null); + + final Schema writerSchema = + AvroSchemaConverter.convertToSchema(physicalDataType.getLogicalType()); + final GenericRecord record = new GenericData.Record(writerSchema); + record.put("user_id", 3L); + record.put("name", "name 3"); + record.put("event_id", 102L); + record.put("payload", "payload 3"); + final ByteArrayOutputStream output = new ByteArrayOutputStream(); + final Encoder encoder; + if (encoding == AvroEncoding.BINARY) { + encoder = EncoderFactory.get().binaryEncoder(output, null); + } else { + encoder = EncoderFactory.get().jsonEncoder(writerSchema, output); + } + new GenericDatumWriter(writerSchema).write(record, encoder); + encoder.flush(); + + final RowData rowData = deserializationSchema.deserialize(output.toByteArray()); + assertThat(rowData.getArity()).isEqualTo(2); + assertThat(rowData.getString(0).toString()).isEqualTo("name 3"); + assertThat(rowData.getLong(1)).isEqualTo(102L); + } + @Test void testOldSeDeNewSchema() { assertThatThrownBy( @@ -188,4 +250,11 @@ private Map getAllOptions(boolean legacyTimestampMapping) { options.put("format", AvroFormatFactory.IDENTIFIER); return options; } + + private Map getAllOptions( + boolean legacyTimestampMapping, AvroEncoding encoding) { + final Map options = getAllOptions(legacyTimestampMapping); + options.put("avro.encoding", encoding.toString()); + return options; + } } diff --git a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroRowDataDeSerializationSchemaTest.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroRowDataDeSerializationSchemaTest.java index 84f9994f6665c0..108f32894fb5fa 100644 --- a/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroRowDataDeSerializationSchemaTest.java +++ b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroRowDataDeSerializationSchemaTest.java @@ -25,6 +25,7 @@ import org.apache.flink.formats.avro.generated.Timestamps; import org.apache.flink.formats.avro.typeutils.AvroSchemaConverter; import org.apache.flink.formats.avro.utils.AvroTestUtils; +import org.apache.flink.table.connector.Projection; import org.apache.flink.table.data.GenericRowData; import org.apache.flink.table.data.RowData; import org.apache.flink.table.data.util.DataFormatConverters; @@ -402,6 +403,103 @@ void testTimestampTypeNewMapping() throws Exception { .isEqualTo("1970-01-01T00:02:03.456"); } + @ParameterizedTest + @EnumSource(AvroEncoding.class) + void testDeserializeWithNonZeroStartingProjection(AvroEncoding encoding) throws Exception { + final DataType physicalDataType = + ROW( + FIELD("user_id", BIGINT()), + FIELD("name", STRING()), + FIELD("event_id", BIGINT()), + FIELD("payload", STRING().notNull())) + .notNull(); + final Schema schema = + AvroSchemaConverter.convertToSchema(physicalDataType.getLogicalType()); + + final GenericRecord record = new GenericData.Record(schema); + record.put(0, 3L); + record.put(1, "name 3"); + record.put(2, 102L); + record.put(3, "payload 3"); + final byte[] input = writeRecord(record, schema, encoding); + + // projection does not start at 0: reads only {name, event_id} + final RowData rowData = + createProjectedDeserializationSchema( + physicalDataType, new int[][] {{1}, {2}}, encoding) + .deserialize(input); + + assertThat(rowData.getArity()).isEqualTo(2); + assertThat(rowData.getString(0).toString()).isEqualTo("name 3"); + assertThat(rowData.getLong(1)).isEqualTo(102L); + } + + @ParameterizedTest + @EnumSource(AvroEncoding.class) + void testDeserializeProjectionIncludingNestedRow(AvroEncoding encoding) throws Exception { + final DataType physicalDataType = + ROW( + FIELD("id", BIGINT()), + FIELD("name", STRING()), + FIELD( + "nested", + ROW(FIELD("a", INT().notNull()), FIELD("b", STRING())) + .notNull())) + .notNull(); + final Schema schema = + AvroSchemaConverter.convertToSchema(physicalDataType.getLogicalType()); + + final GenericRecord nested = new GenericData.Record(schema.getField("nested").schema()); + nested.put(0, 7); + nested.put(1, "inner"); + final GenericRecord record = new GenericData.Record(schema); + record.put(0, 1L); + record.put(1, "outer"); + record.put(2, nested); + final byte[] input = writeRecord(record, schema, encoding); + + // project {name, nested}, i.e. a non-zero starting projection including a nested row + final RowData rowData = + createProjectedDeserializationSchema( + physicalDataType, new int[][] {{1}, {2}}, encoding) + .deserialize(input); + + assertThat(rowData.getArity()).isEqualTo(2); + assertThat(rowData.getString(0).toString()).isEqualTo("outer"); + final RowData nestedOut = rowData.getRow(1, 2); + assertThat(nestedOut.getInt(0)).isEqualTo(7); + assertThat(nestedOut.getString(1).toString()).isEqualTo("inner"); + } + + private static byte[] writeRecord(GenericRecord record, Schema schema, AvroEncoding encoding) + throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + final GenericDatumWriter datumWriter = new GenericDatumWriter<>(schema); + final Encoder encoder = createEncoder(encoding, schema, out); + datumWriter.write(record, encoder); + encoder.flush(); + return out.toByteArray(); + } + + private AvroRowDataDeserializationSchema createProjectedDeserializationSchema( + DataType physicalDataType, int[][] projections, AvroEncoding encoding) + throws Exception { + final DataType producedDataType = Projection.of(projections).project(physicalDataType); + final RowType producedRowType = (RowType) producedDataType.getLogicalType(); + final RowType physicalRowType = (RowType) physicalDataType.getLogicalType(); + + AvroRowDataDeserializationSchema deserializationSchema = + new AvroRowDataDeserializationSchema( + AvroDeserializationSchema.forGeneric( + AvroSchemaConverter.convertToSchema(producedRowType), + AvroSchemaConverter.convertToSchema(physicalRowType), + encoding), + AvroToRowDataConverters.createRowConverter(producedRowType), + InternalTypeInfo.of(producedRowType)); + deserializationSchema.open(null); + return deserializationSchema; + } + private AvroRowDataSerializationSchema createSerializationSchema( DataType dataType, AvroEncoding encoding, boolean legacyTimestampMapping) throws Exception {