Merge pull request #164 from linkedin/xiyan/skip

Add skipReshuffleForParDo option to FlinkPipelineOptions

Author: yanxinyi

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -398,7 +398,7 @@ class BeamModulePlugin implements Plugin<Project> {
 
     // Automatically use the official release version if we are performing a release
     // otherwise append '-SNAPSHOT'
-    project.version = '2.45.53'
+    project.version = '2.45.54'
     if (isLinkedin(project)) {
       project.ext.mavenGroupId = 'com.linkedin.beam'
     }

gradle.properties

@@ -30,8 +30,8 @@ signing.gnupg.useLegacyGpg=true
 # buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy.
 # To build a custom Beam version make sure you change it in both places, see
 # https://github.com/apache/beam/issues/21302.
-version=2.45.53
-sdk_version=2.45.53
+version=2.45.54
+sdk_version=2.45.54
 
 javaVersion=1.8
 

runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java

@@ -305,6 +305,18 @@ public interface FlinkPipelineOptions
 
   void setFlinkConfDir(String confDir);
 
+  @Description(
+      "When true, stateful ParDo operators use DataStreamUtils.reinterpretAsKeyedStream() instead"
+          + " of DataStream.keyBy(), eliminating the HASH network exchange before stateful DoFns."
+          + " Requires that every input PCollection is already correctly key-partitioned by the"
+          + " Flink-compatible hash function (Kafka producer must use a MurmurHash3-aligned"
+          + " partitioner matching Flink's KvToByteBufferKeySelector). Incorrect alignment causes"
+          + " silent state corruption.")
+  @Default.Boolean(false)
+  Boolean getSkipReshuffleForParDo();
+
+  void setSkipReshuffleForParDo(Boolean value);
+
   static FlinkPipelineOptions defaults() {
     return PipelineOptionsFactory.as(FlinkPipelineOptions.class);
   }

runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java

@@ -613,9 +613,14 @@ static <InputT, OutputT> void translateParDo(
             producer != null
                 ? PTransformTranslation.urnForTransformOrNull(context.getProducer(input))
                 : null;
-        // We can skip reshuffle in case previous transform was CPK or GBK
+        // We can skip reshuffle in case previous transform was CPK or GBK,
+        // or if the caller asserts input is already correctly key-partitioned.
         if (PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN.equals(previousUrn)
-            || PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN.equals(previousUrn)) {
+            || PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN.equals(previousUrn)
+            || context
+                .getPipelineOptions()
+                .as(FlinkPipelineOptions.class)
+                .getSkipReshuffleForParDo()) {
           inputDataStream = DataStreamUtils.reinterpretAsKeyedStream(inputDataStream, keySelector);
         } else {
           inputDataStream = inputDataStream.keyBy(keySelector);