Adding skip shuffle shared guard

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.62'
+    project.version = '2.45.63'
     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.62
-sdk_version=2.45.62
+version=2.45.63
+sdk_version=2.45.63
 
 javaVersion=1.8
 

runners/flink/1.17/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java

@@ -56,6 +56,7 @@
 import org.apache.beam.runners.core.TimerInternals.TimerData;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
 import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.runners.flink.FlinkSubtaskContextAware;
 import org.apache.beam.runners.flink.metrics.DoFnRunnerWithMetricsUpdate;
 import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
 import org.apache.beam.runners.flink.metrics.GlobalMetricsUtils;
@@ -495,6 +496,18 @@ public void open() throws Exception {
     // This will be called after initializeState()
     this.doFn = getDoFn();
 
+    // LI-SPECIFIC change: give any DoFn that opts in (by implementing
+    // FlinkSubtaskContextAware) visibility into its own physical placement, so it can
+    // independently detect misrouted/misaligned records (e.g. alongside
+    // SkipShufflePartitionGuard) without needing runner-internal access itself.
+    if (doFn instanceof FlinkSubtaskContextAware) {
+      ((FlinkSubtaskContextAware) doFn)
+          .setFlinkSubtaskContext(
+              getRuntimeContext().getIndexOfThisSubtask(),
+              getRuntimeContext().getNumberOfParallelSubtasks(),
+              getRuntimeContext().getMaxNumberOfParallelSubtasks());
+    }
+
     FlinkPipelineOptions options = serializedOptions.get().as(FlinkPipelineOptions.class);
     doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, options);
 

runners/flink/1.20/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java

@@ -56,6 +56,7 @@
 import org.apache.beam.runners.core.TimerInternals.TimerData;
 import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
 import org.apache.beam.runners.flink.FlinkPipelineOptions;
+import org.apache.beam.runners.flink.FlinkSubtaskContextAware;
 import org.apache.beam.runners.flink.metrics.DoFnRunnerWithMetricsUpdate;
 import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
 import org.apache.beam.runners.flink.metrics.GlobalMetricsUtils;
@@ -495,6 +496,18 @@ public void open() throws Exception {
     // This will be called after initializeState()
     this.doFn = getDoFn();
 
+    // LI-SPECIFIC change: give any DoFn that opts in (by implementing
+    // FlinkSubtaskContextAware) visibility into its own physical placement, so it can
+    // independently detect misrouted/misaligned records (e.g. alongside
+    // SkipShufflePartitionGuard) without needing runner-internal access itself.
+    if (doFn instanceof FlinkSubtaskContextAware) {
+      ((FlinkSubtaskContextAware) doFn)
+          .setFlinkSubtaskContext(
+              getRuntimeContext().getIndexOfThisSubtask(),
+              getRuntimeContext().getNumberOfParallelSubtasks(),
+              getRuntimeContext().getMaxNumberOfParallelSubtasks());
+    }
+
     FlinkPipelineOptions options = serializedOptions.get().as(FlinkPipelineOptions.class);
     doFnInvoker = DoFnInvokers.tryInvokeSetupFor(doFn, options);
 

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

@@ -47,6 +47,7 @@
 import org.apache.beam.runners.core.construction.TransformPayloadTranslatorRegistrar;
 import org.apache.beam.runners.flink.translation.functions.FlinkAssignWindows;
 import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
+import org.apache.beam.runners.flink.translation.utils.SkipShufflePartitionGuard;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.DoFnOperator;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.KvToByteBufferKeySelector;
 import org.apache.beam.runners.flink.translation.wrappers.streaming.SingletonKeyedWorkItem;
@@ -105,6 +106,7 @@
 import org.apache.beam.sdk.values.ValueWithRecordId;
 import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps;
@@ -613,14 +615,56 @@ 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,
-        // 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)
-            || context
-                .getPipelineOptions()
-                .as(FlinkPipelineOptions.class)
-                .getSkipReshuffleForParDo()) {
+        // Flink guarantees producer-verified alignment when the immediately preceding transform
+        // is a CPK or GBK in this same pipeline (its keyBy already partitioned by this exact
+        // key). skipReshuffleForParDo, on the other hand, is an unverified caller assertion about
+        // upstream (e.g. source) partitioning, so unlike the CPK/GBK case it needs a runtime
+        // guard: reinterpretAsKeyedStream performs no validation of its own, and a violated
+        // precondition here silently corrupts keyed state or throws an opaque low-level
+        // KeyGroupRange exception deep in the state backend.
+        boolean producerAlreadyKeyed =
+            PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN.equals(previousUrn)
+                || PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN.equals(previousUrn);
+        boolean skipReshuffleForParDo =
+            context.getPipelineOptions().as(FlinkPipelineOptions.class).getSkipReshuffleForParDo();
+        if (producerAlreadyKeyed || skipReshuffleForParDo) {
+          if (skipReshuffleForParDo && !producerAlreadyKeyed) {
+            int parallelism = context.getExecutionEnvironment().getParallelism();
+            int maxParallelism = context.getExecutionEnvironment().getMaxParallelism();
+            // An unset (-1) maxParallelism is derived by Flink per job vertex and can differ
+            // between the guard and the stateful operator (e.g. after restoring from a
+            // savepoint), so the guard could disagree with the state backend's key-group range.
+            Preconditions.checkState(
+                maxParallelism > 0,
+                "%s requires an explicit maxParallelism (set FlinkPipelineOptions#maxParallelism)"
+                    + " so the runtime partition guard and the stateful operator agree on key"
+                    + " groups.",
+                FlinkPipelineOptions.class.getSimpleName() + "#skipReshuffleForParDo");
+            SkipShufflePartitionGuard.checkParallelismAligned(
+                parallelism,
+                maxParallelism,
+                String.format(
+                    "%s requires parallelism (%s) to be <= maxParallelism (%s), because it relies"
+                        + " on a pointwise reinterpretAsKeyedStream instead of a physical shuffle;"
+                        + " parallelism > maxParallelism means keyBy and reinterpretAsKeyedStream"
+                        + " would route keys differently, and when parallelism < maxParallelism the"
+                        + " upstream source must itself assign contiguous key-group ranges per"
+                        + " subtask (not the default round-robin/striped assignment) for this to be"
+                        + " safe. Either make parallelism <= maxParallelism with a"
+                        + " contiguous-range-aware source partitioner, or disable this option.",
+                    FlinkPipelineOptions.class.getSimpleName() + "#skipReshuffleForParDo",
+                    parallelism,
+                    maxParallelism));
+            inputDataStream =
+                inputDataStream
+                    .flatMap(
+                        SkipShufflePartitionGuard.guard(
+                            (KeySelector<WindowedValue<InputT>, ByteBuffer>) keySelector,
+                            FlinkPipelineOptions.class.getSimpleName() + "#skipReshuffleForParDo",
+                            "this stateful ParDo"))
+                    .returns(inputDataStream.getType())
+                    .name("SkipShufflePartitionGuard");
+          }
           inputDataStream = DataStreamUtils.reinterpretAsKeyedStream(inputDataStream, keySelector);
         } else {
           inputDataStream = inputDataStream.keyBy(keySelector);

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

@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.beam.runners.flink;
+
+import org.apache.beam.sdk.transforms.DoFn;
+
+/**
+ * Opt-in interface for a {@link DoFn} that needs to know which physical Flink subtask it is running
+ * on, and at what parallelism/max-parallelism.
+ *
+ * <p>Beam's SDK deliberately keeps user {@link DoFn}s runner-agnostic, so there is normally no way
+ * for a {@code DoFn} to observe {@code RuntimeContext#getIndexOfThisSubtask()} or the parallelism
+ * it is running with. Some {@code DoFn}s, however, legitimately need this -- for example, a
+ * stateless pre-check {@code DoFn} that verifies (using {@code
+ * org.apache.flink.runtime.state.KeyGroupRangeAssignment#assignKeyToParallelOperator}) that a
+ * record actually belongs on the current subtask before it reaches a downstream keyed operator that
+ * skips the normal shuffle (see {@link
+ * org.apache.beam.runners.flink.translation.utils.SkipShufflePartitionGuard}), so it can reroute
+ * misaligned records (e.g. to a dead-letter output) instead of letting them reach -- and
+ * potentially corrupt -- keyed state.
+ *
+ * <p>The Flink runner's {@code DoFnOperator} calls {@link #setFlinkSubtaskContext} once per
+ * operator instance, during {@code open()}, before any elements are processed. This is
+ * intentionally a minimal, generic hook: it exposes only the three values needed for this kind of
+ * alignment check, and does not expose the full Flink {@code RuntimeContext}, keeping the
+ * runner-specific surface area small.
+ */
+public interface FlinkSubtaskContextAware {
+
+  /**
+   * Called once, before {@code @Setup}/element processing begins, with this operator instance's
+   * physical placement in the current job.
+   *
+   * @param subtaskIndex the 0-based index of the subtask this operator instance is running as
+   * @param parallelism the operator's current parallelism (number of parallel subtasks)
+   * @param maxParallelism the operator's configured max parallelism (number of key groups)
+   */
+  void setFlinkSubtaskContext(int subtaskIndex, int parallelism, int maxParallelism);
+}