Register Beam Schema with JavaSerializer

Author: Abacn

runners/flink/2.0/src/main/java/org/apache/beam/runners/flink/FlinkExecutionEnvironments.java

@@ -36,6 +36,7 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.net.HostAndPort;
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.RuntimeExecutionMode;
+import org.apache.flink.api.common.serialization.SerializerConfigImpl;
 import org.apache.flink.configuration.CheckpointingOptions;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.CoreOptions;
@@ -270,13 +271,7 @@ public static StreamExecutionEnvironment createStreamExecutionEnvironment(
       flinkStreamEnv.getConfig().setAutoWatermarkInterval(options.getAutoWatermarkInterval());
     }
     configureWebUIOptions(flinkStreamEnv.getConfig(), options.as(PipelineOptions.class));
-
-    // Register Kryo serializer for ImmutableBiMap to fix serialization issues in Flink 2
-    ((org.apache.flink.api.common.serialization.SerializerConfigImpl)
-            flinkStreamEnv.getConfig().getSerializerConfig())
-        .registerTypeWithKryoSerializer(
-            org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap.class,
-            org.apache.beam.runners.flink.translation.types.ImmutableMapSerializer.class);
+    configureCustomKryoSerializers(flinkStreamEnv.getConfig());
 
     return flinkStreamEnv;
   }
@@ -301,6 +296,15 @@ private static void configureWebUIOptions(
     }
   }
 
+  // Force Beam schema to use JavaSerializer to fix serialization involving ImmutableMap in Flink 2
+  // TODO:
+  private static void configureCustomKryoSerializers(ExecutionConfig config) {
+    SerializerConfigImpl serializerConfig = (SerializerConfigImpl) config.getSerializerConfig();
+    serializerConfig.registerTypeWithKryoSerializer(
+        org.apache.beam.sdk.schemas.Schema.class,
+        com.esotericsoftware.kryo.serializers.JavaSerializer.class);
+  }
+
   private static class GlobalJobParametersImpl extends ExecutionConfig.GlobalJobParameters {
     private final Map<String, String> jobOptions;