reuvenlax commented on code in PR #38988:
URL: https://github.com/apache/beam/pull/38988#discussion_r3438505208
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java:
##########
@@ -103,6 +104,8 @@ class SimpleParDoFnHelpers<InputT, OutputT, W extends
BoundedWindow> {
// This may additionally be null if it is not a real DoFn but an OldDoFn or
// GroupAlsoByWindowViaWindowSetDoFn
protected @Nullable DoFnSignature fnSignature;
+ boolean activeKey = false;
+ Consumer<K> onStartKey;
Review Comment:
yes
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java:
##########
@@ -299,17 +338,15 @@ void processTimers(
TimerType mode,
DataflowExecutionContext.DataflowStepContext context,
Coder<BoundedWindow> windowCoder,
- Runnable startKey,
Supplier<StreamingSideInputProcessor<?, ?>> sideInputProcessor)
throws Exception {
TimerInternals.TimerData timer = context.getNextFiredTimer(windowCoder);
-
if (timer != null && fnRunner == null) {
// If we need to run reallyStartBundle in here, we need to make sure to
switch the state
// sampler into the start state.
try (Closeable start = operationContext.enterStart()) {
reallyStartBundle();
- startKey.run();
+ this.onStartKey.accept((K) context.stateInternals().getKey());
Review Comment:
done
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/SimpleParDoFnHelpers.java:
##########
@@ -145,6 +149,11 @@ class SimpleParDoFnHelpers<InputT, OutputT, W extends
BoundedWindow> {
this.outputsPerElementTracker = createOutputsPerElementTracker();
this.doFnSchemaInformation = doFnSchemaInformation;
this.sideInputMapping = sideInputMapping;
+ this.onStartKey =
+ k -> {
+ onStartKey.accept(k);
+ this.activeKey = false;
Review Comment:
oops, yes.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java:
##########
@@ -378,15 +394,29 @@ public void processElement(Object elem) throws Exception {
}
@Override
- public void processTimers() {}
+ public void processTimers() throws Exception {
+ if (!activeKey) {
+ onStartKey(null);
Review Comment:
Removed parameter instead, as we never use it.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/PartialGroupByKeyParDoFns.java:
##########
@@ -378,15 +394,29 @@ public void processElement(Object elem) throws Exception {
}
@Override
- public void processTimers() {}
+ public void processTimers() throws Exception {
+ if (!activeKey) {
+ onStartKey(null);
+ }
+ }
@Override
- public void finishKey(Object key) throws Exception {}
+ public void finishKey(Object key) throws Exception {
+ if (!activeKey) {
+ onStartKey((K) key);
+ }
+ sideInputFetcher.persist();
+ sideInputFetcher = null;
+ this.activeKey = false;
+ }
@Override
public void finishBundle() throws Exception {
groupingTable.flush(receiver);
- sideInputFetcher.persist();
+ if (sideInputFetcher != null) {
+ sideInputFetcher.persist();
+ }
+ this.activeKey = false;
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
