[Flink] fix GBK streaming with side input

apache · Nov 6, 2024 · 80a7e22 · 80a7e22
1 parent 687113a
commit 80a7e22
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/.../main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java b/.../main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/DoFnOperator.java
@@ -812,7 +812,9 @@ public final void processElement2(StreamRecord<RawUnionValue> streamRecord) thro
       WindowedValue<InputT> element = it.next();
       // we need to set the correct key in case the operator is
       // a (keyed) window operator
-      setKeyContextElement1(new StreamRecord<>(element));
+      if (keySelector != null) {
+        setCurrentKey(keySelector.getKey(element));
+      }
 
       Iterable<WindowedValue<InputT>> justPushedBack =
           pushbackDoFnRunner.processElementInReadyWindows(element);