This is an automated email from the ASF dual-hosted git repository.
lcwik pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 8964019 [BEAM-2939] Fix splittable DoFn lifecycle. (#11941)
8964019 is described below
commit 8964019b8a5f4c83fb4098c660710e669070b5df
Author: Lukasz Cwik <[email protected]>
AuthorDate: Tue Jun 9 09:07:24 2020 -0700
[BEAM-2939] Fix splittable DoFn lifecycle. (#11941)
* [BEAM-2939] Fix splittable DoFn lifecycle.
getInitialRestriction/splitAndSize should not be wrapped with
startBundle/FinishBundle invocations.
Instead of copying the stateAccessor initialization (used for side inputs)
I made it so that it was initialized only once and cleaned up the
caches/references in the finalizeState call.
* fixup! Address PR comments.
---
.../apache/beam/fn/harness/FnApiDoFnRunner.java | 48 ++++++++++++++++++----
.../beam/fn/harness/state/FnApiStateAccessor.java | 2 +
.../beam/fn/harness/FnApiDoFnRunnerTest.java | 16 ++++----
3 files changed, 51 insertions(+), 15 deletions(-)
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
index 1c8ba32..9735a1c 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
@@ -193,7 +193,21 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
bundleFinalizer);
// Register the appropriate handlers.
- startFunctionRegistry.register(pTransformId, runner::startBundle);
+ switch (pTransform.getSpec().getUrn()) {
+ case PTransformTranslation.PAR_DO_TRANSFORM_URN:
+ case PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN:
+ case
PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
+ startFunctionRegistry.register(pTransformId, runner::startBundle);
+ break;
+ case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
+ // startBundle should not be invoked
+ case PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN:
+ // startBundle should not be invoked
+ case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
+ // startBundle should not be invoked
+ default:
+ // no-op
+ }
String mainInput;
try {
mainInput = ParDoTranslation.getMainInputName(pTransform);
@@ -235,7 +249,21 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
}
pCollectionConsumerRegistry.register(
pTransform.getInputsOrThrow(mainInput), pTransformId,
(FnDataReceiver) mainInputConsumer);
- finishFunctionRegistry.register(pTransformId, runner::finishBundle);
+ switch (pTransform.getSpec().getUrn()) {
+ case PTransformTranslation.PAR_DO_TRANSFORM_URN:
+ case PTransformTranslation.SPLITTABLE_PROCESS_ELEMENTS_URN:
+ case
PTransformTranslation.SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN:
+ finishFunctionRegistry.register(pTransformId, runner::finishBundle);
+ break;
+ case PTransformTranslation.SPLITTABLE_PAIR_WITH_RESTRICTION_URN:
+ // finishBundle should not be invoked
+ case PTransformTranslation.SPLITTABLE_SPLIT_RESTRICTION_URN:
+ // finishBundle should not be invoked
+ case PTransformTranslation.SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN:
+ // finishBundle should not be invoked
+ default:
+ // no-op
+ }
tearDownFunctions.accept(runner::tearDown);
return runner;
}
@@ -269,7 +297,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
private final Collection<FnDataReceiver<WindowedValue<OutputT>>>
mainOutputConsumers;
private final String mainInputId;
- private FnApiStateAccessor<?> stateAccessor;
+ private final FnApiStateAccessor<?> stateAccessor;
private Map<String, BeamFnTimerClient.TimerHandler<?>> timerHandlers;
private final DoFnInvoker<InputT, OutputT> doFnInvoker;
private final StartBundleArgumentProvider startBundleArgumentProvider;
@@ -784,9 +812,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
default:
// no-op
}
- }
- private void startBundle() {
+ // TODO(BEAM-10212): Support caching state data across bundle boundaries.
this.stateAccessor =
new FnApiStateAccessor(
pipelineOptions,
@@ -809,7 +836,9 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
return null;
},
() -> currentWindow);
+ }
+ private void startBundle() {
// Register as a consumer for each timer.
timerHandlers = new HashMap<>();
for (Map.Entry<String, KV<TimeDomain, Coder<Timer<Object>>>>
timerFamilyInfo :
@@ -869,6 +898,9 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
currentWindow = null;
currentRestriction = null;
}
+
+ // TODO(BEAM-10212): Support caching state data across bundle boundaries.
+ this.stateAccessor.finalizeState();
}
private void processElementForSplitRestriction(
@@ -889,6 +921,9 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
currentWatermarkEstimatorState = null;
currentWindow = null;
}
+
+ // TODO(BEAM-10212): Support caching state data across bundle boundaries.
+ this.stateAccessor.finalizeState();
}
/** Internal class to hold the primary and residual roots when converted to
an input element. */
@@ -1164,9 +1199,8 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
doFnInvoker.invokeFinishBundle(finishBundleArgumentProvider);
- // TODO: Support caching state data across bundle boundaries.
+ // TODO(BEAM-10212): Support caching state data across bundle boundaries.
this.stateAccessor.finalizeState();
- this.stateAccessor = null;
}
private void tearDown() {
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
index f516a60..7f4be6b 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java
@@ -459,5 +459,7 @@ public class FnApiStateAccessor<K> implements
SideInputReader, StateBinder {
} catch (Exception e) {
throw new IllegalStateException(e);
}
+ stateFinalizers.clear();
+ stateKeyObjectCache.clear();
}
}
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
index f9eec08..9da16c3 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
@@ -2056,7 +2056,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
null /* bundleSplitListener */,
null /* bundleFinalizer */);
- Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
+ assertTrue(startFunctionRegistry.getFunctions().isEmpty());
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
@@ -2074,7 +2074,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
KV.of("2", KV.of(new OffsetRange(0, 2),
GlobalWindow.TIMESTAMP_MIN_VALUE)))));
mainOutputValues.clear();
- Iterables.getOnlyElement(finishFunctionRegistry.getFunctions()).run();
+ assertTrue(finishFunctionRegistry.getFunctions().isEmpty());
assertThat(mainOutputValues, empty());
Iterables.getOnlyElement(teardownFunctions).run();
@@ -2151,7 +2151,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
null /* bundleSplitListener */,
null /* bundleFinalizer */);
- Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
+ assertTrue(startFunctionRegistry.getFunctions().isEmpty());
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
@@ -2189,7 +2189,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
secondValue.getPane())));
mainOutputValues.clear();
- Iterables.getOnlyElement(finishFunctionRegistry.getFunctions()).run();
+ assertTrue(finishFunctionRegistry.getFunctions().isEmpty());
assertThat(mainOutputValues, empty());
Iterables.getOnlyElement(teardownFunctions).run();
@@ -2265,7 +2265,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
null /* bundleSplitListener */,
null /* bundleFinalizer */);
- Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
+ assertTrue(startFunctionRegistry.getFunctions().isEmpty());
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
@@ -2299,7 +2299,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
1.0))));
mainOutputValues.clear();
- Iterables.getOnlyElement(finishFunctionRegistry.getFunctions()).run();
+ assertTrue(finishFunctionRegistry.getFunctions().isEmpty());
assertThat(mainOutputValues, empty());
Iterables.getOnlyElement(teardownFunctions).run();
@@ -2377,7 +2377,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
null /* bundleSplitListener */,
null /* bundleFinalizer */);
- Iterables.getOnlyElement(startFunctionRegistry.getFunctions()).run();
+ assertTrue(startFunctionRegistry.getFunctions().isEmpty());
mainOutputValues.clear();
assertThat(consumers.keySet(), containsInAnyOrder(inputPCollectionId,
outputPCollectionId));
@@ -2459,7 +2459,7 @@ public class FnApiDoFnRunnerTest implements Serializable {
firstValue.getPane())));
mainOutputValues.clear();
- Iterables.getOnlyElement(finishFunctionRegistry.getFunctions()).run();
+ assertTrue(finishFunctionRegistry.getFunctions().isEmpty());
assertThat(mainOutputValues, empty());
Iterables.getOnlyElement(teardownFunctions).run();
