From 13c012b9d6318c9cd33c3184e24bd80f30859e10 Mon Sep 17 00:00:00 2001 From: Xinyi Yan Date: Mon, 4 May 2026 13:52:13 -0700 Subject: [PATCH] Add skipReshuffleForParDo option to FlinkPipelineOptions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a new boolean option `skipReshuffleForParDo` (default false) to `FlinkPipelineOptions`. When true, `FlinkStreamingTransformTranslators` uses `DataStreamUtils.reinterpretAsKeyedStream()` instead of `.keyBy()` before stateful ParDo operators, eliminating the HASH network exchange. The option extends the existing CPK/GBK skip-reshuffle path — the same `reinterpretAsKeyedStream` branch already used after GroupByKey/Combine is now also taken when this option is set. No new code paths are added; the change is a one-line condition extension. Safety: callers must guarantee that every input PCollection is already correctly key-partitioned using Flink's KvToByteBufferKeySelector hash function. Incorrect alignment causes silent state corruption. Co-Authored-By: Claude Sonnet 4.6 --- .../org/apache/beam/gradle/BeamModulePlugin.groovy | 2 +- gradle.properties | 4 ++-- .../beam/runners/flink/FlinkPipelineOptions.java | 12 ++++++++++++ .../flink/FlinkStreamingTransformTranslators.java | 9 +++++++-- 4 files changed, 22 insertions(+), 5 deletions(-) diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy index e7fc6ac61beb..d181620f97bf 100644 --- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy +++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy @@ -398,7 +398,7 @@ class BeamModulePlugin implements Plugin { // 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' } diff --git a/gradle.properties b/gradle.properties index 46172bc79ead..0f8970f8087f 100644 --- a/gradle.properties +++ b/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 diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java index 829cd559d98a..d0970f75cc24 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java +++ b/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); } diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java index d1224400dfcd..2a69cc673c45 100644 --- a/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java +++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkStreamingTransformTranslators.java @@ -613,9 +613,14 @@ static 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);