[FLINK-34466] Support dataset type facet for Avro

Signed-off-by: Pawel Leszczynski <[email protected]>

Author: pawel-big-lebowski

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/DefaultTypeDatasetFacet.java

@@ -1,9 +1,11 @@
 package org.apache.flink.connector.kafka.lineage;
 
 import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.api.common.serialization.SerializationSchema;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 
 import java.util.Objects;
+import java.util.Optional;
 
 /** Default implementation of {@link KafkaDatasetFacet}. */
 @PublicEvolving
@@ -13,14 +15,26 @@ public class DefaultTypeDatasetFacet implements TypeDatasetFacet {
 
     private final TypeInformation typeInformation;
 
+    private final Optional<SerializationSchema> serializationSchema;
+
     public DefaultTypeDatasetFacet(TypeInformation typeInformation) {
         this.typeInformation = typeInformation;
+        this.serializationSchema = Optional.empty();
+    }
+
+    public DefaultTypeDatasetFacet(SerializationSchema serializationSchema) {
+        this.serializationSchema = Optional.of(serializationSchema);
+        this.typeInformation = null;
     }
 
     public TypeInformation getTypeInformation() {
         return typeInformation;
     }
 
+    public Optional<SerializationSchema> getSerializationSchema() {
+        return serializationSchema;
+    }
+
     public boolean equals(Object o) {
         if (this == o) {
             return true;
@@ -29,12 +43,13 @@ public boolean equals(Object o) {
             return false;
         }
         DefaultTypeDatasetFacet that = (DefaultTypeDatasetFacet) o;
-        return Objects.equals(typeInformation, that.typeInformation);
+        return Objects.equals(typeInformation, that.typeInformation)
+                && Objects.equals(serializationSchema, that.serializationSchema);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(typeInformation);
+        return Objects.hash(typeInformation, serializationSchema);
     }
 
     @Override

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/lineage/TypeDatasetFacet.java

@@ -1,11 +1,26 @@
 package org.apache.flink.connector.kafka.lineage;
 
 import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.api.common.serialization.SerializationSchema;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.streaming.api.lineage.LineageDatasetFacet;
 
+import java.util.Optional;
+
 /** Facet definition to contain type information of source and sink. */
 @PublicEvolving
 public interface TypeDatasetFacet extends LineageDatasetFacet {
     TypeInformation getTypeInformation();
+
+    /**
+     * Sometimes a sink implementing {@link TypeDatasetFacetProvider} is not able to extract type.
+     * This is happening for AvroSerializationSchema due to type erasure problem. In this case, it
+     * makes sense to expose SerializationSchema to the lineage consumer so that it can use it to
+     * extract type information.
+     *
+     * @return
+     */
+    default Optional<SerializationSchema> getSerializationSchema() {
+        return Optional.empty();
+    }
 }

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaRecordSerializationSchemaBuilder.java

@@ -18,6 +18,7 @@
 package org.apache.flink.connector.kafka.sink;
 
 import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.api.common.functions.InvalidTypesException;
 import org.apache.flink.api.common.serialization.SerializationSchema;
 import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
 import org.apache.flink.api.java.typeutils.TypeExtractor;
@@ -432,13 +433,17 @@ public Optional<TypeDatasetFacet> getTypeDatasetFacet() {
                                 ((ResultTypeQueryable<?>) this.valueSerializationSchema)
                                         .getProducedType()));
             } else {
-                // gets type information from serialize method signature
-                Type type =
-                        TypeExtractor.getParameterType(
-                                SerializationSchema.class, valueSerializationSchema.getClass(), 0);
                 try {
+                    Type type =
+                            TypeExtractor.getParameterType(
+                                    SerializationSchema.class,
+                                    valueSerializationSchema.getClass(),
+                                    0);
+
                     return Optional.of(
                             new DefaultTypeDatasetFacet(TypeExtractor.createTypeInfo(type)));
+                } catch (InvalidTypesException e) {
+                    return Optional.of(new DefaultTypeDatasetFacet(valueSerializationSchema));
                 } catch (Exception e) {
                     LOG.info(
                             "Could not extract type information from {}",

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaSink.java

@@ -213,6 +213,10 @@ public LineageVertex getLineageVertex() {
         Optional<TypeDatasetFacet> typeDatasetFacet = Optional.empty();
         if (recordSerializer instanceof TypeDatasetFacetProvider) {
             typeDatasetFacet = ((TypeDatasetFacetProvider) recordSerializer).getTypeDatasetFacet();
+        } else {
+            LOG.info(
+                    "recordSerializer does not implement TypeDatasetFacetProvider: {}",
+                    recordSerializer);
         }
 
         if (typeDatasetFacet.isPresent()) {

flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/DynamicKafkaRecordSerializationSchema.java

@@ -18,6 +18,7 @@
 package org.apache.flink.streaming.connectors.kafka.table;
 
 import org.apache.flink.annotation.Internal;
+import org.apache.flink.api.common.functions.InvalidTypesException;
 import org.apache.flink.api.common.serialization.SerializationSchema;
 import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
 import org.apache.flink.api.java.typeutils.TypeExtractor;
@@ -211,12 +212,14 @@ public Optional<TypeDatasetFacet> getTypeDatasetFacet() {
                     new DefaultTypeDatasetFacet(
                             ((ResultTypeQueryable<?>) this.valueSerialization).getProducedType()));
         } else {
-            // gets type information from serialize method signature
-            Type type =
-                    TypeExtractor.getParameterType(
-                            SerializationSchema.class, valueSerialization.getClass(), 0);
             try {
+                Type type =
+                        TypeExtractor.getParameterType(
+                                SerializationSchema.class, valueSerialization.getClass(), 0);
+
                 return Optional.of(new DefaultTypeDatasetFacet(TypeExtractor.createTypeInfo(type)));
+            } catch (InvalidTypesException e) {
+                return Optional.of(new DefaultTypeDatasetFacet(valueSerialization));
             } catch (Exception e) {
                 LOG.info(
                         "Could not extract type information from {}",

flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaRecordSerializationSchemaBuilderTest.java

@@ -400,6 +400,24 @@ public String apply(Object o) {
                 .isEqualTo(BasicTypeInfo.STRING_TYPE_INFO);
     }
 
+    @Test
+    public void testTypeDatasetFacetWithErasureProblem() {
+        SerializationSchema serializationSchema = element -> new byte[0];
+
+        KafkaRecordSerializationSchema<String> schema =
+                KafkaRecordSerializationSchema.builder()
+                        .setTopic("some-topic")
+                        .setValueSerializationSchema(serializationSchema)
+                        .setKeySerializationSchema(new SimpleStringSchema())
+                        .build();
+
+        assertThat(((TypeDatasetFacetProvider) schema).getTypeDatasetFacet())
+                .isPresent()
+                .get()
+                .extracting(TypeDatasetFacet::getSerializationSchema)
+                .isEqualTo(Optional.of(serializationSchema));
+    }
+
     private static void assertOnlyOneSerializerAllowed(
             List<
                             Function<