From 7a04a229276a4c61c8e96e5923d9ed3b8887c5c9 Mon Sep 17 00:00:00 2001 From: Power John Date: Wed, 15 Jul 2026 23:39:59 +0800 Subject: [PATCH] [FLINK-35324][formats] Avro format can not perform projection pushdown for specific fields --- .../avro/AvroDeserializationSchema.java | 72 +++++++++++++-- .../flink/formats/avro/AvroFormatFactory.java | 16 +++- .../AvroRowDataDeSerializationSchemaTest.java | 92 +++++++++++++++++++ 3 files changed, 169 insertions(+), 11 deletions(-) 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..056f5a566d99af 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 schema 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} + */ + public static AvroDeserializationSchema forGeneric( + Schema schema, @Nullable Schema writer, AvroEncoding encoding) { + return new AvroDeserializationSchema<>(GenericRecord.class, schema, 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); } @@ -236,11 +288,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/AvroRowDataDeSerializationSchemaTest.java b/flink-formats/flink-avro/src/test/java/org/apache/flink/formats/avro/AvroRowDataDeSerializationSchemaTest.java index 84f9994f6665c0..bd53c3fbf3f072 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,97 @@ void testTimestampTypeNewMapping() throws Exception { .isEqualTo("1970-01-01T00:02:03.456"); } + @Test + void testDeserializeWithNonZeroStartingProjection() 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); + + // projection does not start at 0: reads only {name, event_id} + final RowData rowData = + createProjectedDeserializationSchema(physicalDataType, new int[][] {{1}, {2}}) + .deserialize(input); + + assertThat(rowData.getArity()).isEqualTo(2); + assertThat(rowData.getString(0).toString()).isEqualTo("name 3"); + assertThat(rowData.getLong(1)).isEqualTo(102L); + } + + @Test + void testDeserializeProjectionWithNestedRow() 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); + + // project {name, nested}, i.e. a non-zero starting projection including a nested row + final RowData rowData = + createProjectedDeserializationSchema(physicalDataType, new int[][] {{1}, {2}}) + .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) throws Exception { + final ByteArrayOutputStream out = new ByteArrayOutputStream(); + final GenericDatumWriter datumWriter = new GenericDatumWriter<>(schema); + final Encoder encoder = EncoderFactory.get().binaryEncoder(out, null); + datumWriter.write(record, encoder); + encoder.flush(); + return out.toByteArray(); + } + + private AvroRowDataDeserializationSchema createProjectedDeserializationSchema( + DataType physicalDataType, int[][] projections) 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), + AvroEncoding.BINARY), + AvroToRowDataConverters.createRowConverter(producedRowType), + InternalTypeInfo.of(producedRowType)); + deserializationSchema.open(null); + return deserializationSchema; + } + private AvroRowDataSerializationSchema createSerializationSchema( DataType dataType, AvroEncoding encoding, boolean legacyTimestampMapping) throws Exception {