Skip to content

[FLINK-35324][formats] Avro format can not perform projection pushdown for specific fields #28751

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
johntomcat7408-cmyk wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[FLINK-35324][formats] Avro format can not perform projection pushdown for specific fields #28751

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 63 additions & 9 deletions ...ats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroDeserializationSchema.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,7 +71,22 @@ public static AvroDeserializationSchema<GenericRecord> forGeneric(Schema schema)
*/
public static AvroDeserializationSchema<GenericRecord> 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<GenericRecord> forGeneric(
Schema schema, @Nullable Schema writer, AvroEncoding encoding) {
return new AvroDeserializationSchema<>(GenericRecord.class, schema, writer, encoding);
}

/**
Expand Down Expand Up @@ -107,6 +122,9 @@ public static <T extends SpecificRecord> AvroDeserializationSchema<T> 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<T> datumReader;

Expand All @@ -122,6 +140,9 @@ public static <T extends SpecificRecord> AvroDeserializationSchema<T> forSpecifi
/** Avro schema for the reader. */
private transient Schema reader;

/** Avro schema for the writer. */
private transient Schema writer;

/**
* Creates a Avro deserialization schema.
*
Expand All @@ -133,14 +154,31 @@ public static <T extends SpecificRecord> AvroDeserializationSchema<T> forSpecifi
*/
AvroDeserializationSchema(
Class<T> 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<T> 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;
}

Expand All @@ -152,6 +190,11 @@ Schema getReaderSchema() {
return reader;
}

@Nullable
Schema getWriterSchema() {
return writer;
}

MutableByteArrayInputStream getInputStream() {
return inputStream;
}
Expand All @@ -173,9 +216,15 @@ public T deserialize(@Nullable byte[] message) throws IOException {
checkAvroInitialized();
inputStream.setBuffer(message);
Schema readerSchema = getReaderSchema();
Schema writerSchema = getWriterSchema();
GenericDatumReader<T> 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);
Expand All @@ -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);
}
Expand Down Expand Up @@ -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);
}
}
16 changes: 14 additions & 2 deletions flink-formats/flink-avro/src/main/java/org/apache/flink/formats/avro/AvroFormatFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -72,11 +73,22 @@ public DeserializationSchema<RowData> 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<RowData> 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
Expand Down
92 changes: 92 additions & 0 deletions ...vro/src/test/java/org/apache/flink/formats/avro/AvroRowDataDeSerializationSchemaTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<IndexedRecord> 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 {
Expand Down