Adding skip shuffle shared guard

[yanxinyi]

Summary

• Adds SkipShufflePartitionGuard, a runtime safety net for the skipReshuffleForParDo
  optimization (stateful ParDo), which forwards input locally via reinterpretAsKeyedStream
  instead of a real HASH shuffle, relying on the caller's assertion that the input already
  arrives correctly key-partitioned.
• Adds a dropMisalignedRecords mode to the guard: instead of always failing the job fast on
  the first misaligned record, the guard can now log + count
  (skipShufflePartitionGuardDroppedRecords metric) + drop just the offending record and keep
  the job running, controlled by a new skipShuffleGuardDropMisalignedRecords pipeline option.
• Adds FlinkSubtaskContextAware, a minimal opt-in interface that lets any user DoFn learn its
  own (subtaskIndex, parallelism, maxParallelism) — wired into DoFnOperator.open() (1.17 and
  1.20) — so downstream domain code (e.g. CTC's own guard DoFn) can independently detect
  misaligned records and route them to a domain-specific dead-letter sink, which this
  runner-internal guard has no way to do itself.
• Relaxes the graph-construction-time precondition check from parallelism == maxParallelism to
  parallelism <= maxParallelism, matching what reinterpretAsKeyedStream actually requires.
• Scoped to skipReshuffleForParDo only — does not add or touch skipReshuffleForCoGroupByKey.
  CTC is currently the only consumer of this runner's skip-shuffle optimization, and CTC does not
  use CoGroupByKey, so this keeps the change minimal rather than carrying an option/guard path
  nothing exercises yet. (An earlier version of this branch had added CoGroupByKey support
  alongside this refactor; it was dropped before merge to keep the PR scoped to an actual
  consumer's need — see "Notes for reviewers" below.)

Problem

skipReshuffleForParDo skips a real HASH shuffle before a keyed stateful ParDo, relying on the
caller's assertion that the input already arrives partitioned exactly as such a shuffle would
partition it. When that assertion is violated (e.g. an upstream source partitioning bug), the
record silently reaches a stateful operator whose keyed state does not own its key's key-group.
This either:

• throws a cryptic, hard-to-diagnose IllegalArgumentException: Key group X is not in KeyGroupRange{...} deep in the state backend (RocksDBKeyedStateBackend.setCurrentKey), or
• silently corrupts keyed state.

Prior to this PR, there was no runtime validation of this precondition at all, and no way to
recover from (rather than crash on) a misaligned record.

Solution

• SkipShufflePartitionGuard.checkParallelismAligned(...): one-time, graph-construction-time
  check that parallelism <= maxParallelism (required for reinterpretAsKeyedStream to be a valid
  substitute for a real HASH shuffle).
• SkipShufflePartitionGuard.guard(...): a per-record RichFlatMapFunction, inserted just before
  reinterpretAsKeyedStream in the stateful-ParDo translator, that recomputes
  KeyGroupRangeAssignment.assignKeyToParallelOperator(key, maxParallelism, parallelism) and
  compares it to the record's actual subtask. On mismatch: fail-fast (default) or
  log+count+drop (dropMisalignedRecords=true).
• FlinkSubtaskContextAware: a narrow, generic hook (three ints, no full RuntimeContext
  exposure) that any DoFn can opt into, giving domain-specific consumers (like CTC) the same
  runtime facts this guard uses, so they can build their own alignment check with
  domain-specific handling (e.g. dead-lettering) that this runner-internal utility cannot provide.

Files changed

File	Change
FlinkSubtaskContextAware.java (new)	Opt-in interface exposing subtask/parallelism/maxParallelism to a DoFn
SkipShufflePartitionGuard.java (new)	Graph-time + runtime guard for skipReshuffleForParDo, with fail-fast and drop modes
DoFnOperator.java (1.17 and 1.20)	Calls setFlinkSubtaskContext(...) in open() for any DoFn implementing the new interface
FlinkPipelineOptions.java	Adds skipShuffleGuardDropMisalignedRecords pipeline option
FlinkStreamingTransformTranslators.java	Wires the guard into the stateful-ParDo translator; relaxes the parallelism check to <=
ParDoSkipShuffleTest.java, DoFnOperatorTest.java (new/extended)	Tests for the translator's guard behavior (fail-fast + drop mode) and the new FlinkSubtaskContextAware wiring
flink_java_pipeline_options.html, flink_python_pipeline_options.html	Doc-site entry for the new pipeline option
gradle.properties, BeamModulePlugin.groovy	Version bump 2.45.58 → 2.45.59

Testing Done

• :runners:flink:1.20:test — new tests pass (ParDoSkipShuffleTest 6/6,
  DoFnOperatorTest#testFlinkSubtaskContextIsInjectedIntoOpenDoFn); pre-existing unrelated
  JVM/environment failures (JDK 17 module-access) excluded
• :runners:flink:1.20:compileJava/compileTestJava succeed
• :runners:flink:spotlessJavaCheck passes
• Confirmed runners:spark:3:validatesRunnerStreaming's ParDoLifecycleTest failure in CI is
  pre-existing/environment-related, not caused by this diff (no changes touch runners/spark;
  reproduced locally with a real pass once JDK 17 --add-opens flags were forced)

Notes for reviewers

• This is purely additive/internal to the Flink runner — no behavior change unless
  skipReshuffleForParDo is explicitly enabled.
• skipReshuffleForCoGroupByKey was deliberately not included in this PR. CTC is the only
  current consumer of this runner's skip-shuffle feature and doesn't use CoGroupByKey; adding
  guard support for an option nothing exercises yet would be unvalidated, untested-in-production
  scope creep. If/when a real CoGroupByKey consumer needs this, it can be added as its own PR
  against SkipShufflePartitionGuard at that time.
• Downstream consumer: linkedin-multiproduct/ctc PR [BEAM-2247] implement date functions apache/beam#3207 depends on FlinkSubtaskContextAware
  to build a DLQ-routing guard on top of this.

---

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI.

[yananhao12]

Thanks for adding this — failing fast with an actionable message instead of the opaque Key group X is not in KeyGroupRange error is a clear improvement.

Three concerns inline:

1. skipShuffleGuardDropMisalignedRecords can silently drop most of the traffic (e.g. after a rescale). My main ask: remove it or bound it.
2. The guard reads maxParallelism from its own vertex, which can differ from the stateful operator's (side inputs + unset maxP + restore).
3. The drop-mode test would pass even if every record were dropped.