You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by lc...@apache.org on 2018/07/26 18:18:01 UTC
[beam] branch master updated: [BEAM-4658] Update pipeline
representation in runner support libraries to handle timers. (#6050)
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 53d57bf [BEAM-4658] Update pipeline representation in runner support libraries to handle timers. (#6050)
53d57bf is described below
commit 53d57bf642a03d611909015249f3f9b94675b2a2
Author: Lukasz Cwik <lc...@google.com>
AuthorDate: Thu Jul 26 11:17:56 2018 -0700
[BEAM-4658] Update pipeline representation in runner support libraries to handle timers. (#6050)
This is towards supporting timers as PCollections.
---
.../pipeline/src/main/proto/beam_runner_api.proto | 14 ++
.../core/construction/ParDoTranslation.java | 22 +-
.../core/construction/graph/ExecutableStage.java | 39 ++-
.../graph/GreedyPCollectionFusers.java | 80 +++---
.../construction/graph/GreedyPipelineFuser.java | 4 +-
.../core/construction/graph/GreedyStageFuser.java | 6 +-
.../graph/ImmutableExecutableStage.java | 9 +-
.../runners/core/construction/graph/Networks.java | 88 +++++--
.../construction/graph/OutputDeduplicator.java | 1 +
.../core/construction/graph/QueryablePipeline.java | 29 +++
.../construction/graph/SideInputReference.java | 2 +-
.../core/construction/graph/TimerReference.java | 57 +++++
.../construction/graph/UserStateReference.java | 44 +---
.../core/construction/ParDoTranslationTest.java | 3 +-
.../construction/graph/ExecutableStageTest.java | 9 +
.../graph/GreedyPipelineFuserTest.java | 5 +-
.../graph/ImmutableExecutableStageTest.java | 8 +
.../core/construction/graph/NetworksTest.java | 32 +++
.../construction/graph/OutputDeduplicatorTest.java | 8 +
.../portable/RemoteStageEvaluatorFactory.java | 15 +-
.../SplittableRemoteStageEvaluatorFactory.java | 9 +-
.../BatchFlinkExecutableStageContext.java | 4 +-
.../functions/FlinkExecutableStageContext.java | 2 +-
.../functions/FlinkExecutableStageFunction.java | 9 +-
.../streaming/ExecutableStageDoFnOperator.java | 6 +-
.../streaming/ExecutableStageDoFnOperatorTest.java | 32 +--
.../FlinkBatchSideInputHandlerFactoryTest.java | 1 +
.../FlinkExecutableStageFunctionTest.java | 36 +--
runners/java-fn-execution/build.gradle | 1 +
.../control/DockerJobBundleFactory.java | 19 +-
.../fnexecution/control/JobBundleFactory.java | 2 +-
.../control/ProcessBundleDescriptors.java | 162 ++++++++++--
.../runners/fnexecution/control/RemoteBundle.java | 9 +-
.../fnexecution/control/SdkHarnessClient.java | 106 ++++----
.../SingleEnvironmentInstanceJobBundleFactory.java | 29 +--
.../fnexecution/control/StageBundleFactory.java | 4 +-
.../control/DockerJobBundleFactoryTest.java | 4 +-
.../fnexecution/control/RemoteExecutionTest.java | 249 ++++++++++++++++---
.../fnexecution/control/SdkHarnessClientTest.java | 273 +++++++++++----------
39 files changed, 1019 insertions(+), 413 deletions(-)
diff --git a/model/pipeline/src/main/proto/beam_runner_api.proto b/model/pipeline/src/main/proto/beam_runner_api.proto
index 3cec68d..c6ee265 100644
--- a/model/pipeline/src/main/proto/beam_runner_api.proto
+++ b/model/pipeline/src/main/proto/beam_runner_api.proto
@@ -359,6 +359,10 @@ message ExecutableStagePayload {
// this ExecutableStagePayload must be represented within this field.
repeated UserStateId user_states = 7;
+ // The timers required for this executable stage. Each timer of each PTransform within
+ // this ExecutableStagePayload must be represented within this field.
+ repeated TimerId timers = 8;
+
// A reference to a side input. Side inputs are uniquely identified by PTransform id and
// local name.
message SideInputId {
@@ -378,6 +382,16 @@ message ExecutableStagePayload {
// (Required) The local name of this user state for the PTransform that references it.
string local_name = 2;
}
+
+ // A reference to a timer. Timers are uniquely identified by PTransform id and
+ // local name.
+ message TimerId {
+ // (Required) The id of the PTransform that references this timer.
+ string transform_id = 1;
+
+ // (Required) The local name of this timer for the PTransform that references it.
+ string local_name = 2;
+ }
}
// The payload for the primitive ParDo transform.
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
index 26e5656..8e624f7 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/ParDoTranslation.java
@@ -131,9 +131,9 @@ public class ParDoTranslation {
.setPayload(payload.toByteString())
.build());
- String mainInputId = getMainInputId(builder, payload);
+ String mainInputName = getMainInputName(builder, payload);
PCollection<KV<?, ?>> mainInput =
- (PCollection) appliedPTransform.getInputs().get(new TupleTag(mainInputId));
+ (PCollection) appliedPTransform.getInputs().get(new TupleTag(mainInputName));
// https://s.apache.org/beam-portability-timers
// Add a PCollection and coder for each timer. Also treat them as inputs and outputs.
@@ -345,13 +345,23 @@ public class ParDoTranslation {
ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN),
"Unexpected payload type %s",
ptransform.getSpec().getUrn());
- ParDoPayload payload = ParDoPayload.parseFrom(ptransform.getSpec().getPayload());
return components.getPcollectionsOrThrow(
- ptransform.getInputsOrThrow(getMainInputId(ptransform, payload)));
+ ptransform.getInputsOrThrow(getMainInputName(ptransform)));
+ }
+
+ /** Returns the name of the main input of the ptransform. */
+ public static String getMainInputName(RunnerApi.PTransformOrBuilder ptransform)
+ throws IOException {
+ checkArgument(
+ ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN),
+ "Unexpected payload type %s",
+ ptransform.getSpec().getUrn());
+ ParDoPayload payload = ParDoPayload.parseFrom(ptransform.getSpec().getPayload());
+ return getMainInputName(ptransform, payload);
}
- /** Returns the main input id of the ptransform. */
- private static String getMainInputId(
+ /** Returns the name of the main input of the ptransform. */
+ private static String getMainInputName(
RunnerApi.PTransformOrBuilder ptransform, RunnerApi.ParDoPayload payload) {
return Iterables.getOnlyElement(
Sets.difference(
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java
index 486b3b7..3bc83fe 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ExecutableStage.java
@@ -26,6 +26,7 @@ import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.Environment;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.SideInputId;
+import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.TimerId;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.UserStateId;
import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
@@ -72,24 +73,37 @@ public interface ExecutableStage {
/**
* Returns the root {@link PCollectionNode} of this {@link ExecutableStage}. This {@link
* ExecutableStage} executes by reading elements from a Remote gRPC Read Node.
+ *
+ * <p>TODO(BEAM-4658): Add timers as input PCollections to executable stages.
*/
PCollectionNode getInputPCollection();
/**
- * Returns the set of {@link PCollectionNode PCollections} that will be accessed by this {@link
- * ExecutableStage} as side inputs.
+ * Returns a set of descriptors that will be accessed by this {@link ExecutableStage} as side
+ * inputs.
*/
Collection<SideInputReference> getSideInputs();
- /** Returns the set of {@link PTransformNode PTransforms} that contain user state. */
+ /**
+ * Returns the set of descriptors that will consume and produce user state by this {@link
+ * ExecutableStage}.
+ */
Collection<UserStateReference> getUserStates();
/**
+ * Returns the set of descriptors that will consume and produce timers by this {@link
+ * ExecutableStage}.
+ */
+ Collection<TimerReference> getTimers();
+
+ /**
* Returns the leaf {@link PCollectionNode PCollections} of this {@link ExecutableStage}.
*
* <p>All of these {@link PCollectionNode PCollections} are consumed by a {@link PTransformNode
* PTransform} which is not contained within this executable stage, and must be materialized at
* execution time by a Remote gRPC Write Transform.
+ *
+ * <p>TODO(BEAM-4658): Add timers as output PCollections to executable stages.
*/
Collection<PCollectionNode> getOutputPCollections();
@@ -103,8 +117,8 @@ public interface ExecutableStage {
* <ul>
* <li>The {@link PTransform#getSubtransformsList()} is empty. This ensures that executable
* stages are treated as primitive transforms.
- * <li>The only {@link PCollection} in the {@link PTransform#getInputsMap()} is the result of
- * {@link #getInputPCollection()}.
+ * <li>The only {@link PCollection PCollections} in the {@link PTransform#getInputsMap()} is the
+ * result of {@link #getInputPCollection()} and {@link #getSideInputs()}.
* <li>The output {@link PCollection PCollections} in the values of {@link
* PTransform#getOutputsMap()} are the {@link PCollectionNode PCollections} returned by
* {@link #getOutputPCollections()}.
@@ -146,6 +160,13 @@ public interface ExecutableStage {
.setLocalName(userState.localName()));
}
+ for (TimerReference timer : getTimers()) {
+ payload.addTimers(
+ TimerId.newBuilder()
+ .setTransformId(timer.transform().getId())
+ .setLocalName(timer.localName()));
+ }
+
int outputIndex = 0;
for (PCollectionNode output : getOutputPCollections()) {
pt.putOutputs(String.format("materialized_%d", outputIndex), output.getId());
@@ -205,6 +226,12 @@ public interface ExecutableStage {
.stream()
.map(userStateId -> UserStateReference.fromUserStateId(userStateId, components))
.collect(Collectors.toList());
+ List<TimerReference> timers =
+ payload
+ .getTimersList()
+ .stream()
+ .map(timerId -> TimerReference.fromTimerId(timerId, components))
+ .collect(Collectors.toList());
List<PTransformNode> transforms =
payload
.getTransformsList()
@@ -218,6 +245,6 @@ public interface ExecutableStage {
.map(id -> PipelineNode.pCollection(id, components.getPcollectionsOrThrow(id)))
.collect(Collectors.toList());
return ImmutableExecutableStage.of(
- components, environment, input, sideInputs, userStates, transforms, outputs);
+ components, environment, input, sideInputs, userStates, timers, transforms, outputs);
}
}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java
index 64a92f5..3215691 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPCollectionFusers.java
@@ -21,6 +21,7 @@ package org.apache.beam.runners.core.construction.graph;
import static com.google.common.base.Preconditions.checkArgument;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
import java.util.Collection;
import java.util.Map;
import java.util.Optional;
@@ -83,11 +84,12 @@ class GreedyPCollectionFusers {
public static boolean canFuse(
PTransformNode transformNode,
Environment environment,
+ PCollectionNode candidate,
Collection<PCollectionNode> stagePCollections,
QueryablePipeline pipeline) {
return URN_FUSIBILITY_CHECKERS
.getOrDefault(transformNode.getTransform().getSpec().getUrn(), DEFAULT_FUSIBILITY_CHECKER)
- .canFuse(transformNode, environment, stagePCollections, pipeline);
+ .canFuse(transformNode, environment, candidate, stagePCollections, pipeline);
}
public static boolean isCompatible(
@@ -114,6 +116,7 @@ class GreedyPCollectionFusers {
boolean canFuse(
PTransformNode transformNode,
Environment environment,
+ @SuppressWarnings("unused") PCollectionNode candidate,
Collection<PCollectionNode> stagePCollections,
QueryablePipeline pipeline);
}
@@ -137,6 +140,7 @@ class GreedyPCollectionFusers {
private static boolean canFuseParDo(
PTransformNode parDo,
Environment environment,
+ PCollectionNode candidate,
Collection<PCollectionNode> stagePCollections,
QueryablePipeline pipeline) {
Optional<Environment> env = pipeline.getEnvironment(parDo);
@@ -150,41 +154,55 @@ class GreedyPCollectionFusers {
// is never possible.
return false;
}
- if (!pipeline.getSideInputs(parDo).isEmpty()) {
- // At execution time, a Runner is required to only provide inputs to a PTransform that, at
- // the time the PTransform processes them, the associated window is ready in all side inputs
- // that the PTransform consumes. For an arbitrary stage, it is significantly complex for the
- // runner to determine this for each input. As a result, we break fusion to simplify this
- // inspection. In general, a ParDo which consumes side inputs cannot be fused into an
- // executable stage alongside any transforms which are upstream of any of its side inputs.
- return false;
- } else {
- try {
- ParDoPayload payload = ParDoPayload.parseFrom(parDo.getTransform().getSpec().getPayload());
- if (payload.getStateSpecsCount() > 0 || payload.getTimerSpecsCount() > 0) {
- // Inputs to a ParDo that uses State or Timers must be key-partitioned, and elements for
- // a key must execute serially. To avoid checking if the rest of the stage is
- // key-partitioned and preserves keys, these ParDos do not fuse into an existing stage.
- return false;
- }
- } catch (InvalidProtocolBufferException e) {
- throw new IllegalArgumentException(e);
+ try {
+ ParDoPayload payload = ParDoPayload.parseFrom(parDo.getTransform().getSpec().getPayload());
+ if (Maps.filterKeys(
+ parDo.getTransform().getInputsMap(), s -> payload.getTimerSpecsMap().containsKey(s))
+ .values()
+ .contains(candidate.getId())) {
+ // Allow fusion across timer PCollections because they are a self loop.
+ return true;
+ } else if (payload.getStateSpecsCount() > 0 || payload.getTimerSpecsCount() > 0) {
+ // Inputs to a ParDo that uses State or Timers must be key-partitioned, and elements for
+ // a key must execute serially. To avoid checking if the rest of the stage is
+ // key-partitioned and preserves keys, these ParDos do not fuse into an existing stage.
+ return false;
+ } else if (!pipeline.getSideInputs(parDo).isEmpty()) {
+ // At execution time, a Runner is required to only provide inputs to a PTransform that, at
+ // the time the PTransform processes them, the associated window is ready in all side inputs
+ // that the PTransform consumes. For an arbitrary stage, it is significantly complex for the
+ // runner to determine this for each input. As a result, we break fusion to simplify this
+ // inspection. In general, a ParDo which consumes side inputs cannot be fused into an
+ // executable stage alongside any transforms which are upstream of any of its side inputs.
+ return false;
}
+ } catch (InvalidProtocolBufferException e) {
+ throw new IllegalArgumentException(e);
}
return true;
}
private static boolean parDoCompatibility(
PTransformNode parDo, PTransformNode other, QueryablePipeline pipeline) {
- // This is a convenience rather than a strict requirement. In general, a ParDo that consumes
- // side inputs can be fused with other transforms in the same environment which are not
- // upstream of any of the side inputs.
- return pipeline.getSideInputs(parDo).isEmpty()
- // Since we lack the ability to mark upstream transforms as key preserving, we
- // purposefully break fusion here to provide runners the opportunity to insert a
- // grouping operation
- && pipeline.getUserStates(parDo).isEmpty()
- && compatibleEnvironments(parDo, other, pipeline);
+ // Implicitly true if we are attempting to fuse against oneself. This is for timer PCollection which create a loop.
+ return parDo.equals(other)
+ // This is a convenience rather than a strict requirement. In general, a ParDo that consumes
+ // side inputs can be fused with other transforms in the same environment which are not
+ // upstream of any of the side inputs.
+ || (pipeline.getSideInputs(parDo).isEmpty()
+ // We purposefully break fusion here to provide runners the opportunity to insert a
+ // grouping operation to simplify implementing support for ParDo's that contain user state.
+ // We would not need to do this if we had the ability to mark upstream transforms as
+ // key preserving or if runners could execute ParDos containing user state in a distributed
+ // fashion for a single key.
+ && pipeline.getUserStates(parDo).isEmpty()
+ // We purposefully break fusion here to provide runners the opportunity to insert a
+ // grouping operation to simplify implementing support for ParDo's that contain timers.
+ // We would not need to do this if we had the ability to mark upstream transforms as
+ // key preserving or if runners could execute ParDos containing timers in a distributed
+ // fashion for a single key.
+ && pipeline.getTimers(parDo).isEmpty()
+ && compatibleEnvironments(parDo, other, pipeline));
}
/**
@@ -193,6 +211,7 @@ class GreedyPCollectionFusers {
private static boolean canFuseAssignWindows(
PTransformNode window,
Environment environmemnt,
+ @SuppressWarnings("unused") PCollectionNode candidate,
@SuppressWarnings("unused") Collection<PCollectionNode> stagePCollections,
QueryablePipeline pipeline) {
// WindowInto transforms may not have an environment
@@ -256,6 +275,7 @@ class GreedyPCollectionFusers {
private static boolean canAlwaysFuse(
@SuppressWarnings("unused") PTransformNode flatten,
@SuppressWarnings("unused") Environment environment,
+ @SuppressWarnings("unused") PCollectionNode candidate,
@SuppressWarnings("unused") Collection<PCollectionNode> stagePCollections,
@SuppressWarnings("unused") QueryablePipeline pipeline) {
return true;
@@ -264,6 +284,7 @@ class GreedyPCollectionFusers {
private static boolean cannotFuse(
@SuppressWarnings("unused") PTransformNode cannotFuse,
@SuppressWarnings("unused") Environment environment,
+ @SuppressWarnings("unused") PCollectionNode candidate,
@SuppressWarnings("unused") Collection<PCollectionNode> stagePCollections,
@SuppressWarnings("unused") QueryablePipeline pipeline) {
return false;
@@ -284,6 +305,7 @@ class GreedyPCollectionFusers {
private static boolean unknownTransformFusion(
PTransformNode transform,
@SuppressWarnings("unused") Environment environment,
+ @SuppressWarnings("unused") PCollectionNode candidate,
@SuppressWarnings("unused") Collection<PCollectionNode> stagePCollections,
@SuppressWarnings("unused") QueryablePipeline pipeline) {
LOG.debug(
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java
index 8662b0d..945c194 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuser.java
@@ -347,8 +347,8 @@ public class GreedyPipelineFuser {
* PTransformNode} consuming a single materialized {@link PCollectionNode}.
*
* <p>For convenience, {@link CollectionConsumer} implements {@link Comparable}. The natural
- * ordering of {@link CollectionConsumer} is first by the ID of the {@link #consumedCollection()},
- * then by the ID of the {@link #consumingTransform()}.
+ * ordering of {@link CollectionConsumer} is first by the IDs of the {@link
+ * #consumedCollection()}, then by the ID of the {@link #consumingTransform()}.
*/
@AutoValue
abstract static class CollectionConsumer implements Comparable<CollectionConsumer> {
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java
index 89da434..7dfb354 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/GreedyStageFuser.java
@@ -82,6 +82,7 @@ public class GreedyStageFuser {
Set<SideInputReference> sideInputs = new LinkedHashSet<>();
Set<UserStateReference> userStates = new LinkedHashSet<>();
+ Set<TimerReference> timers = new LinkedHashSet<>();
Set<PCollectionNode> fusedCollections = new LinkedHashSet<>();
Set<PCollectionNode> materializedPCollections = new LinkedHashSet<>();
@@ -90,6 +91,7 @@ public class GreedyStageFuser {
fusionCandidates.addAll(pipeline.getOutputPCollections(initialConsumer));
sideInputs.addAll(pipeline.getSideInputs(initialConsumer));
userStates.addAll(pipeline.getUserStates(initialConsumer));
+ timers.addAll(pipeline.getTimers(initialConsumer));
}
while (!fusionCandidates.isEmpty()) {
PCollectionNode candidate = fusionCandidates.poll();
@@ -134,6 +136,7 @@ public class GreedyStageFuser {
inputPCollection,
sideInputs,
userStates,
+ timers,
fusedTransforms.build(),
materializedPCollections);
}
@@ -168,7 +171,8 @@ public class GreedyStageFuser {
Environment environment,
Set<PCollectionNode> fusedPCollections) {
for (PTransformNode node : pipeline.getPerElementConsumers(candidate)) {
- if (!(GreedyPCollectionFusers.canFuse(node, environment, fusedPCollections, pipeline))) {
+ if (!(GreedyPCollectionFusers.canFuse(
+ node, environment, candidate, fusedPCollections, pipeline))) {
// Some of the consumers can't be fused into this subgraph, so the PCollection has to be
// materialized.
// TODO: Potentially, some of the consumers can be fused back into this stage later
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java
index 5b7dedd..f4305e9 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStage.java
@@ -36,6 +36,7 @@ public abstract class ImmutableExecutableStage implements ExecutableStage {
PCollectionNode input,
Collection<SideInputReference> sideInputs,
Collection<UserStateReference> userStates,
+ Collection<TimerReference> timers,
Collection<PTransformNode> transforms,
Collection<PCollectionNode> outputs) {
Components prunedComponents =
@@ -47,7 +48,8 @@ public abstract class ImmutableExecutableStage implements ExecutableStage {
.stream()
.collect(Collectors.toMap(PTransformNode::getId, PTransformNode::getTransform)))
.build();
- return of(prunedComponents, environment, input, sideInputs, userStates, transforms, outputs);
+ return of(
+ prunedComponents, environment, input, sideInputs, userStates, timers, transforms, outputs);
}
public static ImmutableExecutableStage of(
@@ -56,6 +58,7 @@ public abstract class ImmutableExecutableStage implements ExecutableStage {
PCollectionNode input,
Collection<SideInputReference> sideInputs,
Collection<UserStateReference> userStates,
+ Collection<TimerReference> timers,
Collection<PTransformNode> transforms,
Collection<PCollectionNode> outputs) {
return new AutoValue_ImmutableExecutableStage(
@@ -64,6 +67,7 @@ public abstract class ImmutableExecutableStage implements ExecutableStage {
input,
ImmutableSet.copyOf(sideInputs),
ImmutableSet.copyOf(userStates),
+ ImmutableSet.copyOf(timers),
ImmutableSet.copyOf(transforms),
ImmutableSet.copyOf(outputs));
}
@@ -85,6 +89,9 @@ public abstract class ImmutableExecutableStage implements ExecutableStage {
public abstract Collection<UserStateReference> getUserStates();
@Override
+ public abstract Collection<TimerReference> getTimers();
+
+ @Override
public abstract Collection<PTransformNode> getTransforms();
@Override
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/Networks.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/Networks.java
index 9f41a66..efe0da2 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/Networks.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/Networks.java
@@ -23,17 +23,19 @@ import static com.google.common.base.Preconditions.checkNotNull;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Maps;
+import com.google.common.collect.Ordering;
import com.google.common.graph.ElementOrder;
import com.google.common.graph.EndpointPair;
+import com.google.common.graph.Graphs;
import com.google.common.graph.MutableNetwork;
import com.google.common.graph.Network;
import com.google.common.graph.NetworkBuilder;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Comparator;
+import java.util.Deque;
import java.util.HashMap;
import java.util.HashSet;
-import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
@@ -147,16 +149,22 @@ public class Networks {
/**
* Return a set of nodes in sorted topological order.
*
+ * <p>Note that back edges within directed graphs are "broken" returning a topological order for a
+ * directed acyclic network which approximates the original network.
+ *
* <p>Nodes will be considered in the order specified by the {@link Network Network's} {@link
* Network#nodeOrder()}.
*/
public static <NodeT> Iterable<NodeT> topologicalOrder(Network<NodeT, ?> network) {
- return computeTopologicalOrder(network);
+ return computeTopologicalOrder(Graphs.copyOf(network));
}
/**
* Return a set of nodes in sorted topological order.
*
+ * <p>Note that back edges within directed graphs are "broken" returning a topological order for a
+ * directed acyclic network which approximates the original network.
+ *
* <p>Nodes will be considered in the order specified by the {@link
* ElementOrder#sorted(Comparator) sorted ElementOrder} created with the provided comparator.
*/
@@ -183,7 +191,7 @@ public class Networks {
* Network#nodeOrder()}. This ensures that any two Networks with the same nodes and node orders
* produce the same result.
*/
- private static <NodeT> Iterable<NodeT> computeTopologicalOrder(Network<NodeT, ?> network) {
+ private static <NodeT> Iterable<NodeT> computeTopologicalOrder(MutableNetwork<NodeT, ?> network) {
// TODO: (github/guava/2641) Upgrade Guava and remove this method if topological sorting becomes
// supported externally or remove this comment if its not going to be supported externally.
@@ -193,29 +201,65 @@ public class Networks {
"Only networks without self loops are supported, given %s",
network);
- // Linked hashset will prevent duplicates from appearing and will maintain insertion order.
- LinkedHashSet<NodeT> nodes = new LinkedHashSet<>(network.nodes().size());
- Queue<NodeT> processingOrder = new ArrayDeque<>();
- // Add all the roots
- for (NodeT node : network.nodes()) {
- if (network.inDegree(node) == 0) {
- processingOrder.add(node);
- }
- }
+ // Uses the following algorithm:
+ // A FAST & EFFECTIVE HEURISTIC FOR THE FEEDBACK ARC SET PROBLEM
+ // Peter Eades, Xuemin Lin, W. F. Smyth
+ // https://pdfs.semanticscholar.org/c7ed/d9acce96ca357876540e19664eb9d976637f.pdf
+ //
+ // The only edges that are ignored by the algorithm are back edges.
+ // The algorithm (while there are still nodes in the graph):
+ // 1) Removes all sinks from the graph adding them to the beginning of "s2". Continue to do this till there
+ // are no more sinks.
+ // 2) Removes all source from the graph adding them to the end of "s1". Continue to do this till there
+ // are no more sources.
+ // 3) Remote a single node with the highest delta within the graph and add it to the end of "s1".
+ //
+ // The topological order is then the s1 concatenated with s2.
- while (!processingOrder.isEmpty()) {
- NodeT current = processingOrder.remove();
- // If all predecessors have already been added, then we can add this node, otherwise
- // we need to add the node to the back of the processing queue.
- if (nodes.containsAll(network.predecessors(current))) {
- nodes.add(current);
- processingOrder.addAll(network.successors(current));
- } else {
- processingOrder.add(current);
+ Deque<NodeT> s1 = new ArrayDeque<>();
+ Deque<NodeT> s2 = new ArrayDeque<>();
+
+ Ordering<NodeT> maximumOrdering =
+ new Ordering<NodeT>() {
+ @Override
+ public int compare(NodeT t0, NodeT t1) {
+ return (network.outDegree(t0) - network.inDegree(t0))
+ - (network.outDegree(t1) - network.inDegree(t1));
+ }
+ };
+
+ while (!network.nodes().isEmpty()) {
+ boolean nodeRemoved;
+ do {
+ nodeRemoved = false;
+ for (NodeT possibleSink : ImmutableList.copyOf(network.nodes())) {
+ if (network.outDegree(possibleSink) == 0) {
+ network.removeNode(possibleSink);
+ s2.addFirst(possibleSink);
+ nodeRemoved = true;
+ }
+ }
+ } while (nodeRemoved);
+
+ do {
+ nodeRemoved = false;
+ for (NodeT possibleSource : ImmutableList.copyOf(network.nodes())) {
+ if (network.inDegree(possibleSource) == 0) {
+ network.removeNode(possibleSource);
+ s1.addLast(possibleSource);
+ nodeRemoved = true;
+ }
+ }
+ } while (nodeRemoved);
+
+ if (!network.nodes().isEmpty()) {
+ NodeT maximum = maximumOrdering.max(network.nodes());
+ network.removeNode(maximum);
+ s1.addLast(maximum);
}
}
- return nodes;
+ return ImmutableList.<NodeT>builder().addAll(s1).addAll(s2).build();
}
public static <NodeT, EdgeT> String toDot(Network<NodeT, EdgeT> network) {
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java
index 4b06ab6..80396d9 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicator.java
@@ -310,6 +310,7 @@ class OutputDeduplicator {
stage.getInputPCollection(),
stage.getSideInputs(),
stage.getUserStates(),
+ stage.getTimers(),
updatedTransforms,
updatedOutputs);
}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java
index c2c299e..20a3176 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/QueryablePipeline.java
@@ -324,6 +324,23 @@ public class QueryablePipeline {
.collect(Collectors.toSet());
}
+ public Collection<TimerReference> getTimers(PTransformNode transform) {
+ return getLocalTimerNames(transform.getTransform())
+ .stream()
+ .map(
+ localName -> {
+ String transformId = transform.getId();
+ PTransform transformProto = components.getTransformsOrThrow(transformId);
+ String collectionId = transform.getTransform().getInputsOrThrow(localName);
+ PCollection collection = components.getPcollectionsOrThrow(collectionId);
+ return TimerReference.of(
+ PipelineNode.pTransform(transformId, transformProto),
+ localName,
+ PipelineNode.pCollection(collectionId, collection));
+ })
+ .collect(Collectors.toSet());
+ }
+
private Set<String> getLocalSideInputNames(PTransform transform) {
if (PTransformTranslation.PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
try {
@@ -348,6 +365,18 @@ public class QueryablePipeline {
}
}
+ private Set<String> getLocalTimerNames(PTransform transform) {
+ if (PTransformTranslation.PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
+ try {
+ return ParDoPayload.parseFrom(transform.getSpec().getPayload()).getTimerSpecsMap().keySet();
+ } catch (InvalidProtocolBufferException e) {
+ throw new RuntimeException(e);
+ }
+ } else {
+ return Collections.emptySet();
+ }
+ }
+
public Optional<Environment> getEnvironment(PTransformNode parDo) {
return Environments.getEnvironment(parDo.getId(), components);
}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/SideInputReference.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/SideInputReference.java
index b5ca4c7..29e2746 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/SideInputReference.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/SideInputReference.java
@@ -52,7 +52,7 @@ public abstract class SideInputReference {
PipelineNode.pCollection(collectionId, collection));
}
- /** The id of the PTransform that uses this side input. */
+ /** The PTransform that uses this side input. */
public abstract PTransformNode transform();
/** The local name the referencing PTransform uses to refer to this side input. */
public abstract String localName();
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/TimerReference.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/TimerReference.java
new file mode 100644
index 0000000..36dd328
--- /dev/null
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/TimerReference.java
@@ -0,0 +1,57 @@
+/*
+ * 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.core.construction.graph;
+
+import com.google.auto.value.AutoValue;
+import org.apache.beam.model.pipeline.v1.RunnerApi;
+
+/**
+ * A reference to a timer. This includes the PTransform that references the timer as well as the
+ * PCollection referenced. Both are necessary in order to fully resolve a timer.
+ */
+@AutoValue
+public abstract class TimerReference {
+ /** Create a timer reference. */
+ public static TimerReference of(
+ PipelineNode.PTransformNode transform,
+ String localName,
+ PipelineNode.PCollectionNode collection) {
+ return new AutoValue_TimerReference(transform, localName, collection);
+ }
+
+ /** Create a timer reference from a TimerId proto and components. */
+ public static TimerReference fromTimerId(
+ RunnerApi.ExecutableStagePayload.TimerId timerId, RunnerApi.Components components) {
+ String transformId = timerId.getTransformId();
+ String localName = timerId.getLocalName();
+ String collectionId = components.getTransformsOrThrow(transformId).getInputsOrThrow(localName);
+ RunnerApi.PTransform transform = components.getTransformsOrThrow(transformId);
+ RunnerApi.PCollection collection = components.getPcollectionsOrThrow(collectionId);
+ return of(
+ PipelineNode.pTransform(transformId, transform),
+ localName,
+ PipelineNode.pCollection(collectionId, collection));
+ }
+
+ /** The PTransform that uses this timer. */
+ public abstract PipelineNode.PTransformNode transform();
+ /** The local name the referencing PTransform uses to refer to this timer. */
+ public abstract String localName();
+ /** The PCollection that backs this timer. */
+ public abstract PipelineNode.PCollectionNode collection();
+}
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/UserStateReference.java b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/UserStateReference.java
index 08c6a6a..e5e4766 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/UserStateReference.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/graph/UserStateReference.java
@@ -19,19 +19,13 @@
package org.apache.beam.runners.core.construction.graph;
import com.google.auto.value.AutoValue;
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Sets;
-import java.util.Collections;
-import java.util.Set;
+import java.io.IOException;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload.UserStateId;
-import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
-import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
-import org.apache.beam.runners.core.construction.PTransformTranslation;
+import org.apache.beam.runners.core.construction.ParDoTranslation;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
-import org.apache.beam.vendor.protobuf.v3.com.google.protobuf.InvalidProtocolBufferException;
/**
* A reference to user state. This includes the PTransform that references the user state as well as
@@ -49,31 +43,19 @@ public abstract class UserStateReference {
/** Create a user state reference from a UserStateId proto and components. */
public static UserStateReference fromUserStateId(
UserStateId userStateId, RunnerApi.Components components) {
- String transformId = userStateId.getTransformId();
- String localName = userStateId.getLocalName();
-
- PTransform transform = components.getTransformsOrThrow(transformId);
-
- Set<String> sideInputNames = Collections.emptySet();
- if (PTransformTranslation.PAR_DO_TRANSFORM_URN.equals(transform.getSpec().getUrn())) {
- try {
- sideInputNames =
- ParDoPayload.parseFrom(transform.getSpec().getPayload()).getSideInputsMap().keySet();
- } catch (InvalidProtocolBufferException e) {
- throw new RuntimeException(e);
- }
+ PTransform transform = components.getTransformsOrThrow(userStateId.getTransformId());
+ String mainInputCollectionId;
+ try {
+ mainInputCollectionId =
+ transform.getInputsOrThrow(ParDoTranslation.getMainInputName(transform));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
}
-
- // Get the main input PCollection id.
- String collectionId =
- transform.getInputsOrThrow(
- Iterables.getOnlyElement(
- Sets.difference(transform.getInputsMap().keySet(), sideInputNames)));
- PCollection collection = components.getPcollectionsOrThrow(collectionId);
return UserStateReference.of(
- PipelineNode.pTransform(transformId, transform),
- localName,
- PipelineNode.pCollection(collectionId, collection));
+ PipelineNode.pTransform(userStateId.getTransformId(), transform),
+ userStateId.getLocalName(),
+ PipelineNode.pCollection(
+ mainInputCollectionId, components.getPcollectionsOrThrow(mainInputCollectionId)));
}
/** The id of the PTransform that uses this user state. */
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
index 68975dd..aba0a21 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/ParDoTranslationTest.java
@@ -130,7 +130,7 @@ public class ParDoTranslationTest {
@Test
public void toTransformProto() throws Exception {
Map<TupleTag<?>, PValue> inputs = new HashMap<>();
- inputs.put(new TupleTag<KV<Long, String>>() {}, mainInput);
+ inputs.put(new TupleTag<KV<Long, String>>("mainInputName") {}, mainInput);
inputs.putAll(parDo.getAdditionalInputs());
PCollectionTuple output = mainInput.apply(parDo);
@@ -170,6 +170,7 @@ public class ParDoTranslationTest {
assertThat(
ParDoTranslation.getMainInput(protoTransform, components),
equalTo(components.getPcollectionsOrThrow(mainInputId)));
+ assertThat(ParDoTranslation.getMainInputName(protoTransform), equalTo("mainInputName"));
// Validate that the timer PCollections are added correctly.
DoFnSignature signature = DoFnSignatures.signatureForDoFn(parDo.getFn());
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java
index e2c731f..a987993 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ExecutableStageTest.java
@@ -38,6 +38,7 @@ import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload;
import org.apache.beam.model.pipeline.v1.RunnerApi.SdkFunctionSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.SideInput;
import org.apache.beam.model.pipeline.v1.RunnerApi.StateSpec;
+import org.apache.beam.model.pipeline.v1.RunnerApi.TimerSpec;
import org.apache.beam.model.pipeline.v1.RunnerApi.WindowIntoPayload;
import org.apache.beam.runners.core.construction.PTransformTranslation;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
@@ -55,7 +56,9 @@ public class ExecutableStageTest {
PTransform.newBuilder()
.putInputs("input", "input.out")
.putInputs("side_input", "sideInput.in")
+ .putInputs("timer", "timer.out")
.putOutputs("output", "output.out")
+ .putOutputs("timer", "timer.out")
.setSpec(
FunctionSpec.newBuilder()
.setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
@@ -64,11 +67,13 @@ public class ExecutableStageTest {
.setDoFn(SdkFunctionSpec.newBuilder().setEnvironmentId("foo"))
.putSideInputs("side_input", SideInput.getDefaultInstance())
.putStateSpecs("user_state", StateSpec.getDefaultInstance())
+ .putTimerSpecs("timer", TimerSpec.getDefaultInstance())
.build()
.toByteString()))
.build();
PCollection input = PCollection.newBuilder().setUniqueName("input.out").build();
PCollection sideInput = PCollection.newBuilder().setUniqueName("sideInput.in").build();
+ PCollection timer = PCollection.newBuilder().setUniqueName("timer.out").build();
PCollection output = PCollection.newBuilder().setUniqueName("output.out").build();
Components components =
@@ -76,6 +81,7 @@ public class ExecutableStageTest {
.putTransforms("pt", pt)
.putPcollections("input.out", input)
.putPcollections("sideInput.in", sideInput)
+ .putPcollections("timer.out", timer)
.putPcollections("output.out", output)
.putEnvironments("foo", env)
.build();
@@ -87,6 +93,8 @@ public class ExecutableStageTest {
UserStateReference userStateRef =
UserStateReference.of(
transformNode, "user_state", PipelineNode.pCollection("input.out", input));
+ TimerReference timerRef =
+ TimerReference.of(transformNode, "timer", PipelineNode.pCollection("timer.out", timer));
ImmutableExecutableStage stage =
ImmutableExecutableStage.of(
components,
@@ -94,6 +102,7 @@ public class ExecutableStageTest {
PipelineNode.pCollection("input.out", input),
Collections.singleton(sideInputRef),
Collections.singleton(userStateRef),
+ Collections.singleton(timerRef),
Collections.singleton(PipelineNode.pTransform("pt", pt)),
Collections.singleton(PipelineNode.pCollection("output.out", output)));
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuserTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuserTest.java
index 538d3e3..04b7107 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuserTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/GreedyPipelineFuserTest.java
@@ -1005,7 +1005,9 @@ public class GreedyPipelineFuserTest {
PTransform.newBuilder()
.setUniqueName("TimerParDo")
.putInputs("input", "parDo.out")
- .putOutputs("output", "timer.out")
+ .putInputs("timer", "timer.out")
+ .putOutputs("timer", "timer.out")
+ .putOutputs("output", "output.out")
.setSpec(
FunctionSpec.newBuilder()
.setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
@@ -1024,6 +1026,7 @@ public class GreedyPipelineFuserTest {
.putPcollections("parDo.out", pc("parDo.out"))
.putTransforms("timer", timerTransform)
.putPcollections("timer.out", pc("timer.out"))
+ .putPcollections("output.out", pc("output.out"))
.putEnvironments("common", Environment.newBuilder().setUrl("common").build())
.build();
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java
index 61b3cea..4920e17 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/ImmutableExecutableStageTest.java
@@ -49,7 +49,9 @@ public class ImmutableExecutableStageTest {
PTransform.newBuilder()
.putInputs("input", "input.out")
.putInputs("side_input", "sideInput.in")
+ .putInputs("timer", "timer.out")
.putOutputs("output", "output.out")
+ .putOutputs("timer", "timer.out")
.setSpec(
FunctionSpec.newBuilder()
.setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
@@ -58,11 +60,13 @@ public class ImmutableExecutableStageTest {
.setDoFn(RunnerApi.SdkFunctionSpec.newBuilder().setEnvironmentId("foo"))
.putSideInputs("side_input", RunnerApi.SideInput.getDefaultInstance())
.putStateSpecs("user_state", RunnerApi.StateSpec.getDefaultInstance())
+ .putTimerSpecs("timer", RunnerApi.TimerSpec.getDefaultInstance())
.build()
.toByteString()))
.build();
PCollection input = PCollection.newBuilder().setUniqueName("input.out").build();
PCollection sideInput = PCollection.newBuilder().setUniqueName("sideInput.in").build();
+ PCollection timer = PCollection.newBuilder().setUniqueName("timer.out").build();
PCollection output = PCollection.newBuilder().setUniqueName("output.out").build();
Components components =
@@ -71,6 +75,7 @@ public class ImmutableExecutableStageTest {
.putTransforms("other_pt", PTransform.newBuilder().setUniqueName("other").build())
.putPcollections("input.out", input)
.putPcollections("sideInput.in", sideInput)
+ .putPcollections("timer.out", timer)
.putPcollections("output.out", output)
.putEnvironments("foo", env)
.build();
@@ -82,6 +87,8 @@ public class ImmutableExecutableStageTest {
UserStateReference userStateRef =
UserStateReference.of(
transformNode, "user_state", PipelineNode.pCollection("input.out", input));
+ TimerReference timerRef =
+ TimerReference.of(transformNode, "timer", PipelineNode.pCollection("timer.out", timer));
ImmutableExecutableStage stage =
ImmutableExecutableStage.ofFullComponents(
components,
@@ -89,6 +96,7 @@ public class ImmutableExecutableStageTest {
PipelineNode.pCollection("input.out", input),
Collections.singleton(sideInputRef),
Collections.singleton(userStateRef),
+ Collections.singleton(timerRef),
Collections.singleton(PipelineNode.pTransform("pt", pt)),
Collections.singleton(PipelineNode.pCollection("output.out", output)));
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/NetworksTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/NetworksTest.java
index 2dd8834..c42d3a0 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/NetworksTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/NetworksTest.java
@@ -133,6 +133,38 @@ public class NetworksTest {
}
@Test
+ public void testTopologicalOrderWithFeedbackArcs() throws Exception {
+ MutableNetwork<String, String> network = createNetwork();
+ network.addEdge("F", "B", "FB");
+ Iterable<String> sortedNodes = Networks.topologicalOrder(network);
+ Map<String, Integer> nodeToPosition = new HashMap<>(Iterables.size(sortedNodes));
+ int i = 0;
+ for (String node : sortedNodes) {
+ nodeToPosition.put(node, i);
+ i += 1;
+ }
+
+ for (String node : network.nodes()) {
+ for (String descendant :
+ Sets.difference(
+ Networks.reachableNodes(network, ImmutableSet.of(node), Collections.emptySet()),
+ Sets.newHashSet(node))) {
+ if (!Networks.reachableNodes(network, ImmutableSet.of(descendant), Collections.emptySet())
+ .contains(node)) {
+ // We only reliably compare nodes outside of a loop.
+ assertThat(
+ String.format(
+ "Expected position of node %s to be before descendant %s,"
+ + " order returned %s for network %s",
+ node, descendant, sortedNodes, network),
+ nodeToPosition.get(descendant),
+ greaterThan(nodeToPosition.get(node)));
+ }
+ }
+ }
+ }
+
+ @Test
public void testReachableNodesWithEmptyNetwork() {
assertThat(
Networks.reachableNodes(
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java
index 69a625b..098b29a 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/graph/OutputDeduplicatorTest.java
@@ -97,6 +97,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(PipelineNode.pTransform("one", one)),
ImmutableList.of(PipelineNode.pCollection(oneOut.getUniqueName(), oneOut)));
ExecutableStage twoStage =
@@ -106,6 +107,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(PipelineNode.pTransform("two", two)),
ImmutableList.of(PipelineNode.pCollection(twoOut.getUniqueName(), twoOut)));
PTransformNode redTransform = PipelineNode.pTransform("red", red);
@@ -189,6 +191,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(
PipelineNode.pTransform("one", one), PipelineNode.pTransform("shared", shared)),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)));
@@ -199,6 +202,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(
PipelineNode.pTransform("two", two), PipelineNode.pTransform("shared", shared)),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)));
@@ -302,6 +306,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(PipelineNode.pTransform("one", one), sharedTransform),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)));
PTransformNode redTransform = PipelineNode.pTransform("red", red);
@@ -442,6 +447,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(
PipelineNode.pTransform("multi", three),
PipelineNode.pTransform("shared", shared),
@@ -456,6 +462,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(
PipelineNode.pTransform("one", one), PipelineNode.pTransform("shared", shared)),
ImmutableList.of(PipelineNode.pCollection(sharedOut.getUniqueName(), sharedOut)));
@@ -466,6 +473,7 @@ public class OutputDeduplicatorTest {
PipelineNode.pCollection(redOut.getUniqueName(), redOut),
ImmutableList.of(),
ImmutableList.of(),
+ ImmutableList.of(),
ImmutableList.of(
PipelineNode.pTransform("two", two),
PipelineNode.pTransform("otherShared", otherShared)),
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/RemoteStageEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/RemoteStageEvaluatorFactory.java
index ca20949..35525fd 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/RemoteStageEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/RemoteStageEvaluatorFactory.java
@@ -18,6 +18,7 @@
package org.apache.beam.runners.direct.portable;
+import com.google.common.collect.Iterables;
import java.util.ArrayList;
import java.util.Collection;
import javax.annotation.Nullable;
@@ -29,6 +30,7 @@ import org.apache.beam.runners.fnexecution.control.JobBundleFactory;
import org.apache.beam.runners.fnexecution.control.RemoteBundle;
import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.util.WindowedValue;
/**
@@ -59,7 +61,8 @@ class RemoteStageEvaluatorFactory implements TransformEvaluatorFactory {
private class RemoteStageEvaluator<T> implements TransformEvaluator<T> {
private final PTransformNode transform;
- private final RemoteBundle<T> bundle;
+ private final RemoteBundle bundle;
+ private final FnDataReceiver<WindowedValue<?>> mainInput;
private final Collection<UncommittedBundle<?>> outputs;
private RemoteStageEvaluator(PTransformNode transform) throws Exception {
@@ -67,19 +70,21 @@ class RemoteStageEvaluatorFactory implements TransformEvaluatorFactory {
ExecutableStage stage =
ExecutableStage.fromPayload(
ExecutableStagePayload.parseFrom(transform.getTransform().getSpec().getPayload()));
- outputs = new ArrayList<>();
- StageBundleFactory<T> stageFactory = jobFactory.forStage(stage);
- bundle =
+ this.outputs = new ArrayList<>();
+ StageBundleFactory stageFactory = jobFactory.forStage(stage);
+ this.bundle =
stageFactory.getBundle(
BundleFactoryOutputReceiverFactory.create(
bundleFactory, stage.getComponents(), outputs::add),
StateRequestHandler.unsupported(),
BundleProgressHandler.unsupported());
+ // TODO(BEAM-4680): Add support for timers as inputs to the ULR
+ this.mainInput = Iterables.getOnlyElement(bundle.getInputReceivers().values());
}
@Override
public void processElement(WindowedValue<T> element) throws Exception {
- bundle.getInputReceiver().accept(element);
+ mainInput.accept(element);
}
@Override
diff --git a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/SplittableRemoteStageEvaluatorFactory.java b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/SplittableRemoteStageEvaluatorFactory.java
index d1aaf17..ae81b9a 100644
--- a/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/SplittableRemoteStageEvaluatorFactory.java
+++ b/runners/direct-java/src/main/java/org/apache/beam/runners/direct/portable/SplittableRemoteStageEvaluatorFactory.java
@@ -36,6 +36,7 @@ import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
import org.apache.beam.runners.fnexecution.wire.WireCoders;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
@@ -89,7 +90,8 @@ class SplittableRemoteStageEvaluatorFactory implements TransformEvaluatorFactory
private final CopyOnAccessInMemoryStateInternals<byte[]> stateInternals;
private final DirectTimerInternals timerInternals;
- private final RemoteBundle<KV<InputT, RestrictionT>> bundle;
+ private final RemoteBundle bundle;
+ private final FnDataReceiver<WindowedValue<?>> mainInput;
private final Collection<UncommittedBundle<?>> outputs;
private final SDFFeederViaStateAndTimers<InputT, RestrictionT> feeder;
@@ -124,7 +126,7 @@ class SplittableRemoteStageEvaluatorFactory implements TransformEvaluatorFactory
this.bundle =
jobBundleFactory
- .<KV<InputT, RestrictionT>>forStage(stage)
+ .forStage(stage)
.getBundle(
BundleFactoryOutputReceiverFactory.create(
bundleFactory, stage.getComponents(), outputs::add),
@@ -144,6 +146,7 @@ class SplittableRemoteStageEvaluatorFactory implements TransformEvaluatorFactory
}
}
});
+ this.mainInput = Iterables.getOnlyElement(bundle.getInputReceivers().values());
}
@Override
@@ -158,7 +161,7 @@ class SplittableRemoteStageEvaluatorFactory implements TransformEvaluatorFactory
} else {
elementRestriction = feeder.resume(Iterables.getOnlyElement(kwi.timersIterable()));
}
- bundle.getInputReceiver().accept(elementRestriction);
+ mainInput.accept(elementRestriction);
}
@Override
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/BatchFlinkExecutableStageContext.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/BatchFlinkExecutableStageContext.java
index fc43d1c..edbfd70 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/BatchFlinkExecutableStageContext.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/BatchFlinkExecutableStageContext.java
@@ -51,8 +51,8 @@ class BatchFlinkExecutableStageContext implements FlinkExecutableStageContext {
}
@Override
- public <InputT> StageBundleFactory getStageBundleFactory(ExecutableStage executableStage) {
- return jobBundleFactory.<InputT>forStage(executableStage);
+ public StageBundleFactory getStageBundleFactory(ExecutableStage executableStage) {
+ return jobBundleFactory.forStage(executableStage);
}
@Override
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContext.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContext.java
index 1fb5fa9..39c3183 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContext.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageContext.java
@@ -39,7 +39,7 @@ public interface FlinkExecutableStageContext {
return BatchFlinkExecutableStageContext.BatchFactory.INSTANCE;
}
- <InputT> StageBundleFactory<InputT> getStageBundleFactory(ExecutableStage executableStage);
+ StageBundleFactory getStageBundleFactory(ExecutableStage executableStage);
StateRequestHandler getStateRequestHandler(
ExecutableStage executableStage, RuntimeContext runtimeContext);
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
index daf5ab4..3da58e9 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunction.java
@@ -20,6 +20,7 @@ package org.apache.beam.runners.flink.translation.functions;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
+import com.google.common.collect.Iterables;
import java.util.Map;
import javax.annotation.concurrent.GuardedBy;
import org.apache.beam.model.pipeline.v1.RunnerApi;
@@ -66,7 +67,7 @@ public class FlinkExecutableStageFunction<InputT>
private transient RuntimeContext runtimeContext;
private transient FlinkExecutableStageContext stageContext;
private transient StateRequestHandler stateRequestHandler;
- private transient StageBundleFactory<InputT> stageBundleFactory;
+ private transient StageBundleFactory stageBundleFactory;
private transient BundleProgressHandler progressHandler;
public FlinkExecutableStageFunction(
@@ -109,10 +110,12 @@ public class FlinkExecutableStageFunction<InputT>
checkState(
stateRequestHandler != null, "%s not yet prepared", StateRequestHandler.class.getName());
- try (RemoteBundle<InputT> bundle =
+ try (RemoteBundle bundle =
stageBundleFactory.getBundle(
new ReceiverFactory(collector, outputMap), stateRequestHandler, progressHandler)) {
- FnDataReceiver<WindowedValue<InputT>> receiver = bundle.getInputReceiver();
+ // TODO(BEAM-4681): Add support to Flink to support portable timers.
+ FnDataReceiver<WindowedValue<?>> receiver =
+ Iterables.getOnlyElement(bundle.getInputReceivers().values());
for (WindowedValue<InputT> input : iterable) {
receiver.accept(input);
}
diff --git a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
index 7b9c759..2c042c7 100644
--- a/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
+++ b/runners/flink/src/main/java/org/apache/beam/runners/flink/translation/wrappers/streaming/ExecutableStageDoFnOperator.java
@@ -19,6 +19,7 @@ package org.apache.beam.runners.flink.translation.wrappers.streaming;
import static org.apache.flink.util.Preconditions.checkState;
+import com.google.common.collect.Iterables;
import java.util.Collection;
import java.util.List;
import java.util.Map;
@@ -136,11 +137,12 @@ public class ExecutableStageDoFnOperator<InputT, OutputT> extends DoFnOperator<I
checkState(
stateRequestHandler != null, "%s not yet prepared", StateRequestHandler.class.getName());
- try (RemoteBundle<InputT> bundle =
+ try (RemoteBundle bundle =
stageBundleFactory.getBundle(
new ReceiverFactory(outputManager, outputMap), stateRequestHandler, progressHandler)) {
logger.debug(String.format("Sending value: %s", element));
- bundle.getInputReceiver().accept(element);
+ // TODO(BEAM-4681): Add support to Flink to support portable timers.
+ Iterables.getOnlyElement(bundle.getInputReceivers().values()).accept(element);
}
}
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/ExecutableStageDoFnOperatorTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/ExecutableStageDoFnOperatorTest.java
index 8860dd2..dba1c05 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/ExecutableStageDoFnOperatorTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/streaming/ExecutableStageDoFnOperatorTest.java
@@ -118,12 +118,12 @@ public class ExecutableStageDoFnOperatorTest {
testHarness.open();
@SuppressWarnings("unchecked")
- RemoteBundle<Integer> bundle = Mockito.mock(RemoteBundle.class);
+ RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);
@SuppressWarnings("unchecked")
- FnDataReceiver<WindowedValue<Integer>> receiver = Mockito.mock(FnDataReceiver.class);
- when(bundle.getInputReceiver()).thenReturn(receiver);
+ FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
+ when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("pCollectionId", receiver));
Exception expected = new Exception();
doThrow(expected).when(bundle).close();
@@ -141,12 +141,12 @@ public class ExecutableStageDoFnOperatorTest {
getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);
@SuppressWarnings("unchecked")
- RemoteBundle<Integer> bundle = Mockito.mock(RemoteBundle.class);
+ RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);
@SuppressWarnings("unchecked")
- FnDataReceiver<WindowedValue<Integer>> receiver = Mockito.mock(FnDataReceiver.class);
- when(bundle.getInputReceiver()).thenReturn(receiver);
+ FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
+ when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("pCollectionId", receiver));
WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
@@ -206,24 +206,26 @@ public class ExecutableStageDoFnOperatorTest {
WindowedValue<Integer> five = WindowedValue.valueInGlobalWindow(5);
// We use a real StageBundleFactory here in order to exercise the output receiver factory.
- StageBundleFactory<Void> stageBundleFactory =
- new StageBundleFactory<Void>() {
+ StageBundleFactory stageBundleFactory =
+ new StageBundleFactory() {
@Override
- public RemoteBundle<Void> getBundle(
+ public RemoteBundle getBundle(
OutputReceiverFactory receiverFactory,
StateRequestHandler stateRequestHandler,
BundleProgressHandler progressHandler) {
- return new RemoteBundle<Void>() {
+ return new RemoteBundle() {
@Override
public String getId() {
return "bundle-id";
}
@Override
- public FnDataReceiver<WindowedValue<Void>> getInputReceiver() {
- return input -> {
- /* Ignore input*/
- };
+ public Map<String, FnDataReceiver<WindowedValue<?>>> getInputReceivers() {
+ return ImmutableMap.of(
+ "pCollectionId",
+ input -> {
+ /* Ignore input*/
+ });
}
@Override
@@ -240,7 +242,7 @@ public class ExecutableStageDoFnOperatorTest {
public void close() {}
};
// Wire the stage bundle factory into our context.
- when(stageContext.<Void>getStageBundleFactory(any())).thenReturn(stageBundleFactory);
+ when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);
ExecutableStageDoFnOperator<Integer, Integer> operator =
getOperator(
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkBatchSideInputHandlerFactoryTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkBatchSideInputHandlerFactoryTest.java
index c36409c..897289f 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkBatchSideInputHandlerFactoryTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkBatchSideInputHandlerFactoryTest.java
@@ -240,6 +240,7 @@ public class FlinkBatchSideInputHandlerFactoryTest {
sideInputs,
Collections.emptyList(),
Collections.emptyList(),
+ Collections.emptyList(),
Collections.emptyList());
}
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java
index 545789f..1f261f5 100644
--- a/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java
+++ b/runners/flink/src/test/java/org/apache/beam/runners/flink/translation/functions/FlinkExecutableStageFunctionTest.java
@@ -65,7 +65,7 @@ public class FlinkExecutableStageFunctionTest {
@Mock private DistributedCache distributedCache;
@Mock private Collector<RawUnionValue> collector;
@Mock private FlinkExecutableStageContext stageContext;
- @Mock private StageBundleFactory<Integer> stageBundleFactory;
+ @Mock private StageBundleFactory stageBundleFactory;
@Mock private StateRequestHandler stateRequestHandler;
// NOTE: ExecutableStage.fromPayload expects exactly one input, so we provide one here. These unit
@@ -86,7 +86,7 @@ public class FlinkExecutableStageFunctionTest {
MockitoAnnotations.initMocks(this);
when(runtimeContext.getDistributedCache()).thenReturn(distributedCache);
when(stageContext.getStateRequestHandler(any(), any())).thenReturn(stateRequestHandler);
- when(stageContext.<Integer>getStageBundleFactory(any())).thenReturn(stageBundleFactory);
+ when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);
}
@Test
@@ -95,12 +95,12 @@ public class FlinkExecutableStageFunctionTest {
function.open(new Configuration());
@SuppressWarnings("unchecked")
- RemoteBundle<Integer> bundle = Mockito.mock(RemoteBundle.class);
+ RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);
@SuppressWarnings("unchecked")
- FnDataReceiver<WindowedValue<Integer>> receiver = Mockito.mock(FnDataReceiver.class);
- when(bundle.getInputReceiver()).thenReturn(receiver);
+ FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
+ when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("pCollectionId", receiver));
Exception expected = new Exception();
doThrow(expected).when(bundle).close();
@@ -124,12 +124,12 @@ public class FlinkExecutableStageFunctionTest {
function.open(new Configuration());
@SuppressWarnings("unchecked")
- RemoteBundle<Integer> bundle = Mockito.mock(RemoteBundle.class);
+ RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);
@SuppressWarnings("unchecked")
- FnDataReceiver<WindowedValue<Integer>> receiver = Mockito.mock(FnDataReceiver.class);
- when(bundle.getInputReceiver()).thenReturn(receiver);
+ FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
+ when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("pCollectionId", receiver));
WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
@@ -154,24 +154,26 @@ public class FlinkExecutableStageFunctionTest {
"three", 3);
// We use a real StageBundleFactory here in order to exercise the output receiver factory.
- StageBundleFactory<Integer> stageBundleFactory =
- new StageBundleFactory<Integer>() {
+ StageBundleFactory stageBundleFactory =
+ new StageBundleFactory() {
@Override
- public RemoteBundle<Integer> getBundle(
+ public RemoteBundle getBundle(
OutputReceiverFactory receiverFactory,
StateRequestHandler stateRequestHandler,
BundleProgressHandler progressHandler) {
- return new RemoteBundle<Integer>() {
+ return new RemoteBundle() {
@Override
public String getId() {
return "bundle-id";
}
@Override
- public FnDataReceiver<WindowedValue<Integer>> getInputReceiver() {
- return input -> {
- /* Ignore input*/
- };
+ public Map<String, FnDataReceiver<WindowedValue<?>>> getInputReceivers() {
+ return ImmutableMap.of(
+ "pCollectionId",
+ input -> {
+ /* Ignore input*/
+ });
}
@Override
@@ -188,7 +190,7 @@ public class FlinkExecutableStageFunctionTest {
public void close() throws Exception {}
};
// Wire the stage bundle factory into our context.
- when(stageContext.<Integer>getStageBundleFactory(any())).thenReturn(stageBundleFactory);
+ when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);
FlinkExecutableStageFunction<Integer> function = getFunction(outputTagMap);
function.open(new Configuration());
diff --git a/runners/java-fn-execution/build.gradle b/runners/java-fn-execution/build.gradle
index 1979d33..d4a25fa 100644
--- a/runners/java-fn-execution/build.gradle
+++ b/runners/java-fn-execution/build.gradle
@@ -33,6 +33,7 @@ dependencies {
shadow library.java.slf4j_api
testCompile project(":beam-sdks-java-harness")
testCompile project(path: ":beam-runners-core-construction-java", configuration: "shadow")
+ testCompile project(path: ":beam-sdks-java-core", configuration: "shadowTest")
testCompile library.java.junit
testCompile library.java.hamcrest_core
testCompile library.java.hamcrest_library
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactory.java
index 95735d7..fa6f188 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactory.java
@@ -42,7 +42,6 @@ import org.apache.beam.runners.fnexecution.artifact.BeamFileSystemArtifactRetrie
import org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors.ExecutableProcessBundleDescriptor;
import org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor;
import org.apache.beam.runners.fnexecution.data.GrpcDataService;
-import org.apache.beam.runners.fnexecution.data.RemoteInputDestination;
import org.apache.beam.runners.fnexecution.environment.DockerEnvironmentFactory;
import org.apache.beam.runners.fnexecution.environment.EnvironmentFactory;
import org.apache.beam.runners.fnexecution.environment.RemoteEnvironment;
@@ -178,7 +177,7 @@ public class DockerJobBundleFactory implements JobBundleFactory {
}
@Override
- public <T> StageBundleFactory<T> forStage(ExecutableStage executableStage) {
+ public StageBundleFactory forStage(ExecutableStage executableStage) {
WrappedSdkHarnessClient wrappedClient =
environmentCache.getUnchecked(executableStage.getEnvironment());
ExecutableProcessBundleDescriptor processBundleDescriptor;
@@ -238,31 +237,31 @@ public class DockerJobBundleFactory implements JobBundleFactory {
return Platform.OTHER;
}
- private static class SimpleStageBundleFactory<InputT> implements StageBundleFactory<InputT> {
+ private static class SimpleStageBundleFactory implements StageBundleFactory {
- private final BundleProcessor<InputT> processor;
+ private final BundleProcessor processor;
private final ExecutableProcessBundleDescriptor processBundleDescriptor;
// Store the wrapped client in order to keep a live reference into the cache.
private WrappedSdkHarnessClient wrappedClient;
- static <InputT> SimpleStageBundleFactory<InputT> create(
+ static SimpleStageBundleFactory create(
WrappedSdkHarnessClient wrappedClient,
ExecutableProcessBundleDescriptor processBundleDescriptor) {
@SuppressWarnings("unchecked")
- BundleProcessor<InputT> processor =
+ BundleProcessor processor =
wrappedClient
.getClient()
.getProcessor(
processBundleDescriptor.getProcessBundleDescriptor(),
- (RemoteInputDestination) processBundleDescriptor.getRemoteInputDestination(),
+ processBundleDescriptor.getRemoteInputDestinations(),
wrappedClient.getStateServer().getService());
- return new SimpleStageBundleFactory<>(processBundleDescriptor, processor, wrappedClient);
+ return new SimpleStageBundleFactory(processBundleDescriptor, processor, wrappedClient);
}
SimpleStageBundleFactory(
ExecutableProcessBundleDescriptor processBundleDescriptor,
- BundleProcessor<InputT> processor,
+ BundleProcessor processor,
WrappedSdkHarnessClient wrappedClient) {
this.processBundleDescriptor = processBundleDescriptor;
this.processor = processor;
@@ -270,7 +269,7 @@ public class DockerJobBundleFactory implements JobBundleFactory {
}
@Override
- public RemoteBundle<InputT> getBundle(
+ public RemoteBundle getBundle(
OutputReceiverFactory outputReceiverFactory,
StateRequestHandler stateRequestHandler,
BundleProgressHandler progressHandler)
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/JobBundleFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/JobBundleFactory.java
index 11a6825..a892b9c 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/JobBundleFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/JobBundleFactory.java
@@ -27,5 +27,5 @@ import org.apache.beam.runners.core.construction.graph.ExecutableStage;
* <p>Releases all job-scoped resources when closed.
*/
public interface JobBundleFactory extends AutoCloseable {
- <T> StageBundleFactory<T> forStage(ExecutableStage executableStage);
+ StageBundleFactory forStage(ExecutableStage executableStage);
}
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java
index 0338d3e..bc4bdc4 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/ProcessBundleDescriptors.java
@@ -22,6 +22,7 @@ import static com.google.common.base.Preconditions.checkState;
import static org.apache.beam.runners.core.construction.SyntheticComponents.uniqueId;
import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableTable;
import com.google.common.collect.Iterables;
import java.io.IOException;
@@ -40,10 +41,13 @@ import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
+import org.apache.beam.runners.core.construction.SyntheticComponents;
import org.apache.beam.runners.core.construction.graph.ExecutableStage;
+import org.apache.beam.runners.core.construction.graph.PipelineNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PCollectionNode;
import org.apache.beam.runners.core.construction.graph.PipelineNode.PTransformNode;
import org.apache.beam.runners.core.construction.graph.SideInputReference;
+import org.apache.beam.runners.core.construction.graph.TimerReference;
import org.apache.beam.runners.core.construction.graph.UserStateReference;
import org.apache.beam.runners.fnexecution.data.RemoteInputDestination;
import org.apache.beam.runners.fnexecution.wire.LengthPrefixUnknownCoders;
@@ -51,6 +55,9 @@ import org.apache.beam.runners.fnexecution.wire.WireCoders;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.data.RemoteGrpcPortRead;
import org.apache.beam.sdk.fn.data.RemoteGrpcPortWrite;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.util.WindowedValue.FullWindowedValueCoder;
@@ -111,15 +118,28 @@ public class ProcessBundleDescriptors {
Components.Builder components =
stage.getComponents().toBuilder().clearTransforms().putAllTransforms(stageTransforms);
- // The order of these 3 does not matter.
- RemoteInputDestination<WindowedValue<?>> inputDestination =
- addStageInput(dataEndpoint, stage.getInputPCollection(), components);
+ ImmutableMap.Builder<String, RemoteInputDestination<WindowedValue<?>>>
+ inputDestinationsBuilder = ImmutableMap.builder();
+ ImmutableMap.Builder<Target, Coder<WindowedValue<?>>> outputTargetCodersBuilder =
+ ImmutableMap.builder();
- Map<Target, Coder<WindowedValue<?>>> outputTargetCoders =
- addStageOutputs(dataEndpoint, stage.getOutputPCollections(), components);
+ // The order of these does not matter.
+ inputDestinationsBuilder.put(
+ stage.getInputPCollection().getId(),
+ addStageInput(dataEndpoint, stage.getInputPCollection(), components));
+
+ outputTargetCodersBuilder.putAll(
+ addStageOutputs(dataEndpoint, stage.getOutputPCollections(), components));
Map<String, Map<String, SideInputSpec>> sideInputSpecs = addSideInputs(stage, components);
+ Map<String, Map<String, BagUserStateSpec>> bagUserStateSpecs =
+ forBagUserStates(stage, components.build());
+
+ Map<String, Map<String, TimerSpec>> timerSpecs =
+ forTimerSpecs(
+ dataEndpoint, stage, components, inputDestinationsBuilder, outputTargetCodersBuilder);
+
// Copy data from components to ProcessBundleDescriptor.
ProcessBundleDescriptor.Builder bundleDescriptorBuilder =
ProcessBundleDescriptor.newBuilder().setId(id);
@@ -133,15 +153,13 @@ public class ProcessBundleDescriptors {
.putAllWindowingStrategies(components.getWindowingStrategiesMap())
.putAllTransforms(components.getTransformsMap());
- Map<String, Map<String, BagUserStateSpec>> bagUserStateSpecs =
- forBagUserStates(stage, components.build());
-
return ExecutableProcessBundleDescriptor.of(
bundleDescriptorBuilder.build(),
- inputDestination,
- outputTargetCoders,
+ inputDestinationsBuilder.build(),
+ outputTargetCodersBuilder.build(),
sideInputSpecs,
- bagUserStateSpecs);
+ bagUserStateSpecs,
+ timerSpecs);
}
private static Map<Target, Coder<WindowedValue<?>>> addStageOutputs(
@@ -287,6 +305,74 @@ public class ProcessBundleDescriptors {
return idsToSpec.build().rowMap();
}
+ private static Map<String, Map<String, TimerSpec>> forTimerSpecs(
+ ApiServiceDescriptor dataEndpoint,
+ ExecutableStage stage,
+ Components.Builder components,
+ ImmutableMap.Builder<String, RemoteInputDestination<WindowedValue<?>>> remoteInputsBuilder,
+ ImmutableMap.Builder<Target, Coder<WindowedValue<?>>> outputTargetCodersBuilder)
+ throws IOException {
+ ImmutableTable.Builder<String, String, TimerSpec> idsToSpec = ImmutableTable.builder();
+ for (TimerReference timerReference : stage.getTimers()) {
+ RunnerApi.ParDoPayload payload =
+ RunnerApi.ParDoPayload.parseFrom(
+ timerReference.transform().getTransform().getSpec().getPayload());
+ RunnerApi.TimeDomain.Enum timeDomain =
+ payload.getTimerSpecsOrThrow(timerReference.localName()).getTimeDomain();
+ org.apache.beam.sdk.state.TimerSpec spec;
+ switch (timeDomain) {
+ case EVENT_TIME:
+ spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+ break;
+ case PROCESSING_TIME:
+ spec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+ break;
+ case SYNCHRONIZED_PROCESSING_TIME:
+ spec = TimerSpecs.timer(TimeDomain.SYNCHRONIZED_PROCESSING_TIME);
+ break;
+ default:
+ throw new IllegalArgumentException(String.format("Unknown time domain %s", timeDomain));
+ }
+
+ remoteInputsBuilder.put(
+ timerReference.collection().getId(),
+ addStageInput(dataEndpoint, timerReference.collection(), components));
+ // "Unroll" the timer PCollection to make the execution tree a DAG.
+ String outputTimerPCollectionId =
+ SyntheticComponents.uniqueId(
+ String.format("%s.out", timerReference.collection().getId()),
+ components.getPcollectionsMap()::containsKey);
+ components.putPcollections(
+ outputTimerPCollectionId, timerReference.collection().getPCollection());
+ TargetEncoding targetEncoding =
+ addStageOutput(
+ dataEndpoint,
+ components,
+ PipelineNode.pCollection(
+ outputTimerPCollectionId, timerReference.collection().getPCollection()));
+ outputTargetCodersBuilder.put(targetEncoding.getTarget(), targetEncoding.getCoder());
+ components.putTransforms(
+ timerReference.transform().getId(),
+ // Since a transform can have more then one timer, update the transform inside components and not the original
+ components
+ .getTransformsOrThrow(timerReference.transform().getId())
+ .toBuilder()
+ .putOutputs(timerReference.localName(), outputTimerPCollectionId)
+ .build());
+
+ idsToSpec.put(
+ timerReference.transform().getId(),
+ timerReference.localName(),
+ TimerSpec.of(
+ timerReference.transform().getId(),
+ timerReference.localName(),
+ timerReference.collection().getId(),
+ targetEncoding.getTarget(),
+ spec));
+ }
+ return idsToSpec.build().rowMap();
+ }
+
@AutoValue
abstract static class TargetEncoding {
abstract BeamFnApi.Target getTarget();
@@ -348,15 +434,43 @@ public class ProcessBundleDescriptors {
public abstract Coder<W> windowCoder();
}
+ /**
+ * A container type storing references to the key, timer and payload coders and the remote input
+ * destination used when handling timer requests.
+ */
+ @AutoValue
+ public abstract static class TimerSpec<K, V, W extends BoundedWindow> {
+ static <K, V, W extends BoundedWindow> TimerSpec<K, V, W> of(
+ String transformId,
+ String timerId,
+ String collectionId,
+ Target outputTarget,
+ org.apache.beam.sdk.state.TimerSpec timerSpec) {
+ return new AutoValue_ProcessBundleDescriptors_TimerSpec(
+ transformId, timerId, collectionId, outputTarget, timerSpec);
+ }
+
+ public abstract String transformId();
+
+ public abstract String timerId();
+
+ public abstract String collectionId();
+
+ public abstract Target outputTarget();
+
+ public abstract org.apache.beam.sdk.state.TimerSpec getTimerSpec();
+ }
+
/** */
@AutoValue
public abstract static class ExecutableProcessBundleDescriptor {
public static ExecutableProcessBundleDescriptor of(
ProcessBundleDescriptor descriptor,
- RemoteInputDestination<WindowedValue<?>> inputDestination,
+ Map<String, RemoteInputDestination<WindowedValue<?>>> inputDestinations,
Map<BeamFnApi.Target, Coder<WindowedValue<?>>> outputTargetCoders,
Map<String, Map<String, SideInputSpec>> sideInputSpecs,
- Map<String, Map<String, BagUserStateSpec>> bagUserStateSpecs) {
+ Map<String, Map<String, BagUserStateSpec>> bagUserStateSpecs,
+ Map<String, Map<String, TimerSpec>> timerSpecs) {
ImmutableTable.Builder copyOfSideInputSpecs = ImmutableTable.builder();
for (Map.Entry<String, Map<String, SideInputSpec>> outer : sideInputSpecs.entrySet()) {
for (Map.Entry<String, SideInputSpec> inner : outer.getValue().entrySet()) {
@@ -369,21 +483,29 @@ public class ProcessBundleDescriptors {
copyOfBagUserStateSpecs.put(outer.getKey(), inner.getKey(), inner.getValue());
}
}
+ ImmutableTable.Builder copyOfTimerSpecs = ImmutableTable.builder();
+ for (Map.Entry<String, Map<String, TimerSpec>> outer : timerSpecs.entrySet()) {
+ for (Map.Entry<String, TimerSpec> inner : outer.getValue().entrySet()) {
+ copyOfTimerSpecs.put(outer.getKey(), inner.getKey(), inner.getValue());
+ }
+ }
return new AutoValue_ProcessBundleDescriptors_ExecutableProcessBundleDescriptor(
descriptor,
- inputDestination,
+ inputDestinations,
Collections.unmodifiableMap(outputTargetCoders),
copyOfSideInputSpecs.build().rowMap(),
- copyOfBagUserStateSpecs.build().rowMap());
+ copyOfBagUserStateSpecs.build().rowMap(),
+ copyOfTimerSpecs.build().rowMap());
}
public abstract ProcessBundleDescriptor getProcessBundleDescriptor();
/**
- * Get the {@link RemoteInputDestination} that input data are sent to the {@link
+ * Get {@link RemoteInputDestination}s that input data/timers are sent to the {@link
* ProcessBundleDescriptor} over.
*/
- public abstract RemoteInputDestination<WindowedValue<?>> getRemoteInputDestination();
+ public abstract Map<String, RemoteInputDestination<WindowedValue<?>>>
+ getRemoteInputDestinations();
/**
* Get all of the targets materialized by this {@link ExecutableProcessBundleDescriptor} and the
@@ -402,5 +524,11 @@ public class ProcessBundleDescriptors {
* states} that are used during execution.
*/
public abstract Map<String, Map<String, BagUserStateSpec>> getBagUserStateSpecs();
+
+ /**
+ * Get a mapping from PTransform id to timer id to {@link TimerSpec timer specs} that are used
+ * during execution.
+ */
+ public abstract Map<String, Map<String, TimerSpec>> getTimerSpecs();
}
}
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/RemoteBundle.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/RemoteBundle.java
index d28fdfd..ff3066a 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/RemoteBundle.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/RemoteBundle.java
@@ -18,6 +18,7 @@
package org.apache.beam.runners.fnexecution.control;
+import java.util.Map;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.util.WindowedValue;
@@ -29,15 +30,15 @@ import org.apache.beam.sdk.util.WindowedValue;
* <p>When a RemoteBundle is closed, it will block until bundle processing is finished on remote
* resources, and throw an exception if bundle processing has failed.
*/
-public interface RemoteBundle<InputT> extends AutoCloseable {
+public interface RemoteBundle extends AutoCloseable {
/** Get an id used to represent this bundle. */
String getId();
/**
- * Get a {@link FnDataReceiver receiver} which consumes input elements, forwarding them to the
- * remote environment.
+ * Get a map of PCollection ids to {@link FnDataReceiver receiver}s which consume input elements,
+ * forwarding them to the remote environment.
*/
- FnDataReceiver<WindowedValue<InputT>> getInputReceiver();
+ Map<String, FnDataReceiver<WindowedValue<?>>> getInputReceivers();
/**
* Closes this bundle. This causes the input {@link FnDataReceiver} to be closed (future calls to
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java
index 605a8ea..72fc002 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java
@@ -20,6 +20,7 @@ package org.apache.beam.runners.fnexecution.control;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
+import com.google.common.collect.ImmutableMap;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CompletionStage;
@@ -37,6 +38,7 @@ import org.apache.beam.runners.fnexecution.data.RemoteInputDestination;
import org.apache.beam.runners.fnexecution.state.StateDelegator;
import org.apache.beam.runners.fnexecution.state.StateDelegator.Registration;
import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
+import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.fn.IdGenerator;
import org.apache.beam.sdk.fn.IdGenerators;
import org.apache.beam.sdk.fn.data.CloseableFnDataReceiver;
@@ -60,20 +62,20 @@ public class SdkHarnessClient implements AutoCloseable {
/**
* A processor capable of creating bundles for some registered {@link ProcessBundleDescriptor}.
*/
- public class BundleProcessor<T> {
+ public class BundleProcessor {
private final ProcessBundleDescriptor processBundleDescriptor;
private final CompletionStage<RegisterResponse> registrationFuture;
- private final RemoteInputDestination<WindowedValue<T>> remoteInput;
+ private final Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputs;
private final StateDelegator stateDelegator;
private BundleProcessor(
ProcessBundleDescriptor processBundleDescriptor,
CompletionStage<RegisterResponse> registrationFuture,
- RemoteInputDestination<WindowedValue<T>> remoteInput,
+ Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputs,
StateDelegator stateDelegator) {
this.processBundleDescriptor = processBundleDescriptor;
this.registrationFuture = registrationFuture;
- this.remoteInput = remoteInput;
+ this.remoteInputs = remoteInputs;
this.stateDelegator = stateDelegator;
}
@@ -90,13 +92,14 @@ public class SdkHarnessClient implements AutoCloseable {
* <p>NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
*
* <pre>{@code
- * try (ActiveBundle<InputT> bundle = SdkHarnessClient.newBundle(...)) {
- * FnDataReceiver<InputT> inputReceiver = bundle.getInputReceiver();
- * // send all elements ...
+ * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
+ * FnDataReceiver<InputT> inputReceiver =
+ * (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
+ * // send all main input elements ...
* }
* }</pre>
*/
- public ActiveBundle<T> newBundle(
+ public ActiveBundle newBundle(
Map<BeamFnApi.Target, RemoteOutputReceiver<?>> outputReceivers,
BundleProgressHandler progressHandler) {
return newBundle(
@@ -119,13 +122,14 @@ public class SdkHarnessClient implements AutoCloseable {
* <p>NOTE: It is important to {@link #close()} each bundle after all elements are emitted.
*
* <pre>{@code
- * try (ActiveBundle<InputT> bundle = SdkHarnessClient.newBundle(...)) {
- * FnDataReceiver<InputT> inputReceiver = bundle.getInputReceiver();
+ * try (ActiveBundle bundle = SdkHarnessClient.newBundle(...)) {
+ * FnDataReceiver<InputT> inputReceiver =
+ * (FnDataReceiver) bundle.getInputReceivers().get(mainPCollectionId);
* // send all elements ...
* }
* }</pre>
*/
- public ActiveBundle<T> newBundle(
+ public ActiveBundle newBundle(
Map<BeamFnApi.Target, RemoteOutputReceiver<?>> outputReceivers,
StateRequestHandler stateRequestHandler,
BundleProgressHandler progressHandler) {
@@ -159,14 +163,21 @@ public class SdkHarnessClient implements AutoCloseable {
outputClients.put(targetReceiver.getKey(), outputClient);
}
- CloseableFnDataReceiver<WindowedValue<T>> dataReceiver =
- fnApiDataService.send(
- LogicalEndpoint.of(bundleId, remoteInput.getTarget()), remoteInput.getCoder());
+ ImmutableMap.Builder<String, CloseableFnDataReceiver<WindowedValue<?>>> dataReceiversBuilder =
+ ImmutableMap.builder();
+ for (Map.Entry<String, RemoteInputDestination<WindowedValue<?>>> remoteInput :
+ remoteInputs.entrySet()) {
+ dataReceiversBuilder.put(
+ remoteInput.getKey(),
+ fnApiDataService.send(
+ LogicalEndpoint.of(bundleId, remoteInput.getValue().getTarget()),
+ (Coder) remoteInput.getValue().getCoder()));
+ }
- return new ActiveBundle<>(
+ return new ActiveBundle(
bundleId,
specificResponse,
- dataReceiver,
+ dataReceiversBuilder.build(),
outputClients,
stateDelegator.registerForProcessBundleInstructionId(bundleId, stateRequestHandler),
progressHandler);
@@ -182,10 +193,10 @@ public class SdkHarnessClient implements AutoCloseable {
}
/** An active bundle for a particular {@link BeamFnApi.ProcessBundleDescriptor}. */
- public static class ActiveBundle<InputT> implements RemoteBundle<InputT> {
+ public static class ActiveBundle implements RemoteBundle {
private final String bundleId;
private final CompletionStage<BeamFnApi.ProcessBundleResponse> response;
- private final CloseableFnDataReceiver<WindowedValue<InputT>> inputReceiver;
+ private final Map<String, CloseableFnDataReceiver<WindowedValue<?>>> inputReceivers;
private final Map<BeamFnApi.Target, InboundDataClient> outputClients;
private final StateDelegator.Registration stateRegistration;
private final BundleProgressHandler progressHandler;
@@ -193,13 +204,13 @@ public class SdkHarnessClient implements AutoCloseable {
private ActiveBundle(
String bundleId,
CompletionStage<ProcessBundleResponse> response,
- CloseableFnDataReceiver<WindowedValue<InputT>> inputReceiver,
+ Map<String, CloseableFnDataReceiver<WindowedValue<?>>> inputReceivers,
Map<Target, InboundDataClient> outputClients,
- Registration stateRegistration,
+ StateDelegator.Registration stateRegistration,
BundleProgressHandler progressHandler) {
this.bundleId = bundleId;
this.response = response;
- this.inputReceiver = inputReceiver;
+ this.inputReceivers = inputReceivers;
this.outputClients = outputClients;
this.stateRegistration = stateRegistration;
this.progressHandler = progressHandler;
@@ -212,19 +223,19 @@ public class SdkHarnessClient implements AutoCloseable {
}
/**
- * Returns a {@link FnDataReceiver receiver} which consumes input elements forwarding them to
- * the SDK.
+ * Get a map of PCollection ids to {@link FnDataReceiver receiver}s which consume input
+ * elements, forwarding them to the remote environment.
*/
@Override
- public FnDataReceiver<WindowedValue<InputT>> getInputReceiver() {
- return inputReceiver;
+ public Map<String, FnDataReceiver<WindowedValue<?>>> getInputReceivers() {
+ return (Map) inputReceivers;
}
/**
* Blocks till bundle processing is finished. This is comprised of:
*
* <ul>
- * <li>closing the {@link #getInputReceiver() input receiver}.
+ * <li>closing each {@link #getInputReceivers() input receiver}.
* <li>waiting for the SDK to say that processing the bundle is finished.
* <li>waiting for all inbound data clients to complete
* </ul>
@@ -235,10 +246,16 @@ public class SdkHarnessClient implements AutoCloseable {
@Override
public void close() throws Exception {
Exception exception = null;
- try {
- inputReceiver.close();
- } catch (Exception e) {
- exception = e;
+ for (CloseableFnDataReceiver<?> inputReceiver : inputReceivers.values()) {
+ try {
+ inputReceiver.close();
+ } catch (Exception e) {
+ if (exception == null) {
+ exception = e;
+ } else {
+ exception.addSuppressed(e);
+ }
+ }
}
try {
// We don't have to worry about the completion stage.
@@ -295,7 +312,7 @@ public class SdkHarnessClient implements AutoCloseable {
private final InstructionRequestHandler fnApiControlClient;
private final FnDataService fnApiDataService;
- private final ConcurrentHashMap<String, BundleProcessor<?>> clientProcessors;
+ private final ConcurrentHashMap<String, BundleProcessor> clientProcessors;
private SdkHarnessClient(
InstructionRequestHandler fnApiControlClient,
@@ -336,16 +353,16 @@ public class SdkHarnessClient implements AutoCloseable {
* ProcessBundleDescriptor#getId() process bundle descriptor id}. A previously created instance
* may be returned.
*/
- public <T> BundleProcessor<T> getProcessor(
+ public BundleProcessor getProcessor(
BeamFnApi.ProcessBundleDescriptor descriptor,
- RemoteInputDestination<WindowedValue<T>> remoteInputDesination) {
+ Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputDesinations) {
checkState(
!descriptor.hasStateApiServiceDescriptor(),
"The %s cannot support a %s containing a state %s.",
BundleProcessor.class.getSimpleName(),
BeamFnApi.ProcessBundleDescriptor.class.getSimpleName(),
Endpoints.ApiServiceDescriptor.class.getSimpleName());
- return getProcessor(descriptor, remoteInputDesination, NoOpStateDelegator.INSTANCE);
+ return getProcessor(descriptor, remoteInputDesinations, NoOpStateDelegator.INSTANCE);
}
/**
@@ -362,15 +379,14 @@ public class SdkHarnessClient implements AutoCloseable {
* ProcessBundleDescriptor#getId() process bundle descriptor id}. A previously created instance
* may be returned.
*/
- public <T> BundleProcessor<T> getProcessor(
+ public BundleProcessor getProcessor(
BeamFnApi.ProcessBundleDescriptor descriptor,
- RemoteInputDestination<WindowedValue<T>> remoteInputDesination,
+ Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputDesinations,
StateDelegator stateDelegator) {
@SuppressWarnings("unchecked")
- BundleProcessor<T> bundleProcessor =
- (BundleProcessor<T>)
- clientProcessors.computeIfAbsent(
- descriptor.getId(), s -> create(descriptor, remoteInputDesination, stateDelegator));
+ BundleProcessor bundleProcessor =
+ clientProcessors.computeIfAbsent(
+ descriptor.getId(), s -> create(descriptor, remoteInputDesinations, stateDelegator));
checkArgument(
bundleProcessor.processBundleDescriptor.equals(descriptor),
"The provided %s with id %s collides with an existing %s with the same id but "
@@ -407,9 +423,9 @@ public class SdkHarnessClient implements AutoCloseable {
}
/** Registers a {@link BeamFnApi.ProcessBundleDescriptor} for future processing. */
- private <T> BundleProcessor<T> create(
+ private BundleProcessor create(
BeamFnApi.ProcessBundleDescriptor processBundleDescriptor,
- RemoteInputDestination<WindowedValue<T>> remoteInputDestination,
+ Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputDestinations,
StateDelegator stateDelegator) {
LOG.debug("Registering {}", processBundleDescriptor);
@@ -427,11 +443,11 @@ public class SdkHarnessClient implements AutoCloseable {
CompletionStage<RegisterResponse> registerResponseFuture =
genericResponse.thenApply(InstructionResponse::getRegister);
- BundleProcessor<T> bundleProcessor =
- new BundleProcessor<>(
+ BundleProcessor bundleProcessor =
+ new BundleProcessor(
processBundleDescriptor,
registerResponseFuture,
- remoteInputDestination,
+ remoteInputDestinations,
stateDelegator);
return bundleProcessor;
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SingleEnvironmentInstanceJobBundleFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SingleEnvironmentInstanceJobBundleFactory.java
index 122bf3e..c50c9ad 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SingleEnvironmentInstanceJobBundleFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SingleEnvironmentInstanceJobBundleFactory.java
@@ -30,7 +30,6 @@ import org.apache.beam.runners.core.construction.graph.ExecutableStage;
import org.apache.beam.runners.fnexecution.GrpcFnServer;
import org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors.ExecutableProcessBundleDescriptor;
import org.apache.beam.runners.fnexecution.data.GrpcDataService;
-import org.apache.beam.runners.fnexecution.data.RemoteInputDestination;
import org.apache.beam.runners.fnexecution.environment.EnvironmentFactory;
import org.apache.beam.runners.fnexecution.environment.RemoteEnvironment;
import org.apache.beam.runners.fnexecution.state.GrpcStateService;
@@ -63,7 +62,7 @@ public class SingleEnvironmentInstanceJobBundleFactory implements JobBundleFacto
private final GrpcFnServer<GrpcDataService> dataService;
private final GrpcFnServer<GrpcStateService> stateService;
- private final ConcurrentMap<ExecutableStage, StageBundleFactory<?>> stageBundleFactories =
+ private final ConcurrentMap<ExecutableStage, StageBundleFactory> stageBundleFactories =
new ConcurrentHashMap<>();
private final ConcurrentMap<Environment, RemoteEnvironment> environments =
new ConcurrentHashMap<>();
@@ -80,12 +79,11 @@ public class SingleEnvironmentInstanceJobBundleFactory implements JobBundleFacto
}
@Override
- public <T> StageBundleFactory<T> forStage(ExecutableStage executableStage) {
- return (StageBundleFactory<T>)
- stageBundleFactories.computeIfAbsent(executableStage, this::createBundleFactory);
+ public StageBundleFactory forStage(ExecutableStage executableStage) {
+ return stageBundleFactories.computeIfAbsent(executableStage, this::createBundleFactory);
}
- private <T> StageBundleFactory<T> createBundleFactory(ExecutableStage stage) {
+ private StageBundleFactory createBundleFactory(ExecutableStage stage) {
RemoteEnvironment remoteEnv =
environments.computeIfAbsent(
stage.getEnvironment(),
@@ -108,14 +106,12 @@ public class SingleEnvironmentInstanceJobBundleFactory implements JobBundleFacto
} catch (IOException e) {
throw new RuntimeException(e);
}
- RemoteInputDestination<? super WindowedValue<?>> destination =
- descriptor.getRemoteInputDestination();
- SdkHarnessClient.BundleProcessor<T> bundleProcessor =
+ SdkHarnessClient.BundleProcessor bundleProcessor =
sdkHarnessClient.getProcessor(
descriptor.getProcessBundleDescriptor(),
- (RemoteInputDestination<WindowedValue<T>>) (RemoteInputDestination) destination,
+ descriptor.getRemoteInputDestinations(),
stateService.getService());
- return new BundleProcessorStageBundleFactory<>(descriptor, bundleProcessor);
+ return new BundleProcessorStageBundleFactory(descriptor, bundleProcessor);
}
@Override
@@ -137,19 +133,18 @@ public class SingleEnvironmentInstanceJobBundleFactory implements JobBundleFacto
}
}
- private static class BundleProcessorStageBundleFactory<T> implements StageBundleFactory<T> {
+ private static class BundleProcessorStageBundleFactory implements StageBundleFactory {
private final ExecutableProcessBundleDescriptor descriptor;
- private final SdkHarnessClient.BundleProcessor<T> processor;
+ private final SdkHarnessClient.BundleProcessor processor;
private BundleProcessorStageBundleFactory(
- ExecutableProcessBundleDescriptor descriptor,
- SdkHarnessClient.BundleProcessor<T> processor) {
+ ExecutableProcessBundleDescriptor descriptor, SdkHarnessClient.BundleProcessor processor) {
this.descriptor = descriptor;
this.processor = processor;
}
@Override
- public RemoteBundle<T> getBundle(
+ public RemoteBundle getBundle(
OutputReceiverFactory outputReceiverFactory,
StateRequestHandler stateRequestHandler,
BundleProgressHandler progressHandler) {
@@ -167,7 +162,7 @@ public class SingleEnvironmentInstanceJobBundleFactory implements JobBundleFacto
outputReceiverFactory.create(bundleOutputPCollection);
outputReceivers.put(
targetCoders.getKey(),
- RemoteOutputReceiver.of((Coder) targetCoders.getValue(), outputReceiver));
+ RemoteOutputReceiver.of(targetCoders.getValue(), outputReceiver));
}
return processor.newBundle(outputReceivers, stateRequestHandler, progressHandler);
}
diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/StageBundleFactory.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/StageBundleFactory.java
index 8dae649..3b8b8d2 100644
--- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/StageBundleFactory.java
+++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/StageBundleFactory.java
@@ -28,9 +28,9 @@ import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
* <p>Closing a StageBundleFactory signals that the stage has completed and any resources bound to
* its lifetime can be cleaned up.
*/
-public interface StageBundleFactory<T> extends AutoCloseable {
+public interface StageBundleFactory extends AutoCloseable {
/** Get a new {@link RemoteBundle bundle} for processing the data in an executable stage. */
- RemoteBundle<T> getBundle(
+ RemoteBundle getBundle(
OutputReceiverFactory outputReceiverFactory,
StateRequestHandler stateRequestHandler,
BundleProgressHandler progressHandler)
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactoryTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactoryTest.java
index dbf3dea..3ad361d 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactoryTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/DockerJobBundleFactoryTest.java
@@ -122,8 +122,8 @@ public class DockerJobBundleFactoryTest {
loggingServer,
retrievalServer,
provisioningServer)) {
- StageBundleFactory<?> bf1 = bundleFactory.forStage(getExecutableStage(environment));
- StageBundleFactory<?> bf2 = bundleFactory.forStage(getExecutableStage(environment));
+ StageBundleFactory bf1 = bundleFactory.forStage(getExecutableStage(environment));
+ StageBundleFactory bf2 = bundleFactory.forStage(getExecutableStage(environment));
// NOTE: We hang on to stage bundle references to ensure their underlying environments are not
// garbage collected. For additional safety, we print the factories to ensure the referernces
// are not optimized away.
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
index 13a533a..593390e 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/RemoteExecutionTest.java
@@ -21,14 +21,17 @@ package org.apache.beam.runners.fnexecution.control;
import static com.google.common.base.Preconditions.checkState;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.junit.Assert.assertThat;
+import static org.junit.Assert.fail;
import com.google.common.base.Optional;
+import com.google.common.collect.Collections2;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Iterators;
+import com.google.common.collect.Sets;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import java.io.Serializable;
-import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
@@ -38,6 +41,7 @@ import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
+import java.util.Set;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -57,7 +61,6 @@ import org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors.Exec
import org.apache.beam.runners.fnexecution.control.SdkHarnessClient.ActiveBundle;
import org.apache.beam.runners.fnexecution.control.SdkHarnessClient.BundleProcessor;
import org.apache.beam.runners.fnexecution.data.GrpcDataService;
-import org.apache.beam.runners.fnexecution.data.RemoteInputDestination;
import org.apache.beam.runners.fnexecution.logging.GrpcLoggingService;
import org.apache.beam.runners.fnexecution.logging.Slf4jLogWriter;
import org.apache.beam.runners.fnexecution.state.GrpcStateService;
@@ -69,10 +72,12 @@ import org.apache.beam.runners.fnexecution.state.StateRequestHandlers.SideInputH
import org.apache.beam.runners.fnexecution.state.StateRequestHandlers.SideInputHandlerFactory;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.BigEndianLongCoder;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
+import org.apache.beam.sdk.coders.VoidCoder;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.fn.stream.OutboundObserverFactory;
import org.apache.beam.sdk.fn.test.InProcessManagedChannelFactory;
@@ -81,6 +86,11 @@ import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.ReadableState;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
+import org.apache.beam.sdk.state.TimeDomain;
+import org.apache.beam.sdk.state.Timer;
+import org.apache.beam.sdk.state.TimerSpec;
+import org.apache.beam.sdk.state.TimerSpecs;
+import org.apache.beam.sdk.testing.ResetDateTimeProvider;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.transforms.Impulse;
@@ -88,6 +98,7 @@ import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
@@ -96,8 +107,11 @@ import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.vendor.protobuf.v3.com.google.protobuf.ByteString;
import org.hamcrest.collection.IsEmptyIterable;
import org.hamcrest.collection.IsIterableContainingInOrder;
+import org.joda.time.DateTimeUtils;
+import org.joda.time.Duration;
import org.junit.After;
import org.junit.Before;
+import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@@ -108,6 +122,8 @@ import org.junit.runners.JUnit4;
*/
@RunWith(JUnit4.class)
public class RemoteExecutionTest implements Serializable {
+ @Rule public transient ResetDateTimeProvider resetDateTimeProvider = new ResetDateTimeProvider();
+
private transient GrpcFnServer<FnApiControlClientPoolService> controlServer;
private transient GrpcFnServer<GrpcDataService> dataServer;
private transient GrpcFnServer<GrpcStateService> stateServer;
@@ -161,7 +177,7 @@ public class RemoteExecutionTest implements Serializable {
});
// TODO: https://issues.apache.org/jira/browse/BEAM-4149 Use proper worker id.
InstructionRequestHandler controlClient =
- clientPool.getSource().take("", Duration.ofSeconds(2));
+ clientPool.getSource().take("", java.time.Duration.ofSeconds(2));
this.controlClient = SdkHarnessClient.usingFnApiClient(controlClient, dataServer.getService());
}
@@ -225,12 +241,9 @@ public class RemoteExecutionTest implements Serializable {
ProcessBundleDescriptors.fromExecutableStage(
"my_stage", stage, dataServer.getApiServiceDescriptor());
- @SuppressWarnings({"unchecked", "rawtypes"})
- RemoteInputDestination<WindowedValue<byte[]>> remoteDestination =
- (RemoteInputDestination) descriptor.getRemoteInputDestination();
-
- BundleProcessor<byte[]> processor =
- controlClient.getProcessor(descriptor.getProcessBundleDescriptor(), remoteDestination);
+ BundleProcessor processor =
+ controlClient.getProcessor(
+ descriptor.getProcessBundleDescriptor(), descriptor.getRemoteInputDestinations());
Map<Target, ? super Coder<WindowedValue<?>>> outputTargets = descriptor.getOutputTargetCoders();
Map<Target, Collection<? super WindowedValue<?>>> outputValues = new HashMap<>();
Map<Target, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
@@ -246,9 +259,10 @@ public class RemoteExecutionTest implements Serializable {
}
// The impulse example
- try (ActiveBundle<byte[]> bundle =
+ try (ActiveBundle bundle =
processor.newBundle(outputReceivers, BundleProgressHandler.unsupported())) {
- bundle.getInputReceiver().accept(WindowedValue.valueInGlobalWindow(new byte[0]));
+ Iterables.getOnlyElement(bundle.getInputReceivers().values())
+ .accept(WindowedValue.valueInGlobalWindow(new byte[0]));
}
for (Collection<? super WindowedValue<?>> windowedValues : outputValues.values()) {
@@ -312,13 +326,11 @@ public class RemoteExecutionTest implements Serializable {
dataServer.getApiServiceDescriptor(),
stateServer.getApiServiceDescriptor());
- @SuppressWarnings({"unchecked", "rawtypes"})
- RemoteInputDestination<WindowedValue<byte[]>> remoteDestination =
- (RemoteInputDestination) descriptor.getRemoteInputDestination();
-
- BundleProcessor<byte[]> processor =
+ BundleProcessor processor =
controlClient.getProcessor(
- descriptor.getProcessBundleDescriptor(), remoteDestination, stateDelegator);
+ descriptor.getProcessBundleDescriptor(),
+ descriptor.getRemoteInputDestinations(),
+ stateDelegator);
Map<Target, Coder<WindowedValue<?>>> outputTargets = descriptor.getOutputTargetCoders();
Map<Target, Collection<WindowedValue<?>>> outputValues = new HashMap<>();
Map<Target, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
@@ -361,15 +373,13 @@ public class RemoteExecutionTest implements Serializable {
});
BundleProgressHandler progressHandler = BundleProgressHandler.unsupported();
- try (ActiveBundle<byte[]> bundle =
+ try (ActiveBundle bundle =
processor.newBundle(outputReceivers, stateRequestHandler, progressHandler)) {
- bundle
- .getInputReceiver()
+ Iterables.getOnlyElement(bundle.getInputReceivers().values())
.accept(
WindowedValue.valueInGlobalWindow(
CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "X")));
- bundle
- .getInputReceiver()
+ Iterables.getOnlyElement(bundle.getInputReceivers().values())
.accept(
WindowedValue.valueInGlobalWindow(
CoderUtils.encodeToByteArray(StringUtf8Coder.of(), "Y")));
@@ -447,13 +457,11 @@ public class RemoteExecutionTest implements Serializable {
dataServer.getApiServiceDescriptor(),
stateServer.getApiServiceDescriptor());
- @SuppressWarnings({"unchecked", "rawtypes"})
- RemoteInputDestination<WindowedValue<KV<byte[], byte[]>>> remoteDestination =
- (RemoteInputDestination) descriptor.getRemoteInputDestination();
-
- BundleProcessor<KV<byte[], byte[]>> processor =
+ BundleProcessor processor =
controlClient.getProcessor(
- descriptor.getProcessBundleDescriptor(), remoteDestination, stateDelegator);
+ descriptor.getProcessBundleDescriptor(),
+ descriptor.getRemoteInputDestinations(),
+ stateDelegator);
Map<Target, Coder<WindowedValue<?>>> outputTargets = descriptor.getOutputTargetCoders();
Map<Target, Collection<WindowedValue<?>>> outputValues = new HashMap<>();
Map<Target, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
@@ -515,10 +523,11 @@ public class RemoteExecutionTest implements Serializable {
}
});
- try (ActiveBundle<KV<byte[], byte[]>> bundle =
+ try (ActiveBundle bundle =
processor.newBundle(
outputReceivers, stateRequestHandler, BundleProgressHandler.unsupported())) {
- bundle.getInputReceiver().accept(WindowedValue.valueInGlobalWindow(kvBytes("X", "Y")));
+ Iterables.getOnlyElement(bundle.getInputReceivers().values())
+ .accept(WindowedValue.valueInGlobalWindow(kvBytes("X", "Y")));
}
for (Collection<WindowedValue<?>> windowedValues : outputValues.values()) {
assertThat(
@@ -542,6 +551,164 @@ public class RemoteExecutionTest implements Serializable {
assertThat(userStateData.get(stateId2), IsEmptyIterable.emptyIterable());
}
+ @Test
+ public void testExecutionWithTimer() throws Exception {
+ Pipeline p = Pipeline.create();
+ final String timerId = "foo";
+ final String timerId2 = "foo2";
+
+ p.apply("impulse", Impulse.create())
+ .apply(
+ "create",
+ ParDo.of(
+ new DoFn<byte[], KV<String, String>>() {
+ @ProcessElement
+ public void process(ProcessContext ctxt) {}
+ }))
+ .setCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()))
+ .apply(
+ "timer",
+ ParDo.of(
+ new DoFn<KV<String, String>, KV<String, String>>() {
+ @TimerId("event")
+ private final TimerSpec eventTimerSpec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ @TimerId("processing")
+ private final TimerSpec processingTimerSpec =
+ TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext context,
+ @TimerId("event") Timer eventTimeTimer,
+ @TimerId("processing") Timer processingTimeTimer) {
+ context.output(KV.of("main" + context.element().getKey(), ""));
+ eventTimeTimer.set(context.timestamp().plus(1L));
+ processingTimeTimer.offset(Duration.millis(2L));
+ processingTimeTimer.setRelative();
+ }
+
+ @OnTimer("event")
+ public void eventTimer(
+ OnTimerContext context,
+ @TimerId("event") Timer eventTimeTimer,
+ @TimerId("processing") Timer processingTimeTimer) {
+ context.output(KV.of("event", ""));
+ eventTimeTimer.set(context.timestamp().plus(11L));
+ processingTimeTimer.offset(Duration.millis(12L));
+ processingTimeTimer.setRelative();
+ }
+
+ @OnTimer("processing")
+ public void processingTimer(
+ OnTimerContext context,
+ @TimerId("event") Timer eventTimeTimer,
+ @TimerId("processing") Timer processingTimeTimer) {
+ context.output(KV.of("processing", ""));
+ eventTimeTimer.set(context.timestamp().plus(21L));
+ processingTimeTimer.offset(Duration.millis(22L));
+ processingTimeTimer.setRelative();
+ }
+ }))
+ // Force the output to be materialized
+ .apply("gbk", GroupByKey.create());
+
+ RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);
+ FusedPipeline fused = GreedyPipelineFuser.fuse(pipelineProto);
+ Optional<ExecutableStage> optionalStage =
+ Iterables.tryFind(
+ fused.getFusedStages(), (ExecutableStage stage) -> !stage.getTimers().isEmpty());
+ checkState(optionalStage.isPresent(), "Expected a stage with timers.");
+ ExecutableStage stage = optionalStage.get();
+
+ ExecutableProcessBundleDescriptor descriptor =
+ ProcessBundleDescriptors.fromExecutableStage(
+ "test_stage",
+ stage,
+ dataServer.getApiServiceDescriptor(),
+ stateServer.getApiServiceDescriptor());
+
+ BundleProcessor processor =
+ controlClient.getProcessor(
+ descriptor.getProcessBundleDescriptor(),
+ descriptor.getRemoteInputDestinations(),
+ stateDelegator);
+ Map<Target, Coder<WindowedValue<?>>> outputTargets = descriptor.getOutputTargetCoders();
+ Map<Target, Collection<WindowedValue<?>>> outputValues = new HashMap<>();
+ Map<Target, RemoteOutputReceiver<?>> outputReceivers = new HashMap<>();
+ for (Entry<Target, Coder<WindowedValue<?>>> targetCoder : outputTargets.entrySet()) {
+ List<WindowedValue<?>> outputContents = Collections.synchronizedList(new ArrayList<>());
+ outputValues.put(targetCoder.getKey(), outputContents);
+ outputReceivers.put(
+ targetCoder.getKey(),
+ RemoteOutputReceiver.of(targetCoder.getValue(), outputContents::add));
+ }
+
+ String eventTimeInputPCollectionId = null;
+ Target eventTimeOutputTarget = null;
+ String processingTimeInputPCollectionId = null;
+ Target processingTimeOutputTarget = null;
+ for (Map<String, ProcessBundleDescriptors.TimerSpec> timerSpecs :
+ descriptor.getTimerSpecs().values()) {
+ for (ProcessBundleDescriptors.TimerSpec timerSpec : timerSpecs.values()) {
+ if (TimeDomain.EVENT_TIME.equals(timerSpec.getTimerSpec().getTimeDomain())) {
+ eventTimeInputPCollectionId = timerSpec.collectionId();
+ eventTimeOutputTarget = timerSpec.outputTarget();
+ } else if (TimeDomain.PROCESSING_TIME.equals(timerSpec.getTimerSpec().getTimeDomain())) {
+ processingTimeInputPCollectionId = timerSpec.collectionId();
+ processingTimeOutputTarget = timerSpec.outputTarget();
+ } else {
+ fail(String.format("Unknown timer specification %s", timerSpec));
+ }
+ }
+ }
+
+ // Set the current system time to a fixed value to get stable values for processing time timer output.
+ DateTimeUtils.setCurrentMillisFixed(BoundedWindow.TIMESTAMP_MIN_VALUE.getMillis());
+
+ try (ActiveBundle bundle =
+ processor.newBundle(
+ outputReceivers,
+ StateRequestHandler.unsupported(),
+ BundleProgressHandler.unsupported())) {
+ bundle
+ .getInputReceivers()
+ .get(stage.getInputPCollection().getId())
+ .accept(WindowedValue.valueInGlobalWindow(kvBytes("X", "X")));
+ bundle
+ .getInputReceivers()
+ .get(eventTimeInputPCollectionId)
+ .accept(WindowedValue.valueInGlobalWindow(timerBytes("Y", 100L)));
+ bundle
+ .getInputReceivers()
+ .get(processingTimeInputPCollectionId)
+ .accept(WindowedValue.valueInGlobalWindow(timerBytes("Z", 200L)));
+ }
+ Set<Target> timerOutputTargets =
+ ImmutableSet.of(eventTimeOutputTarget, processingTimeOutputTarget);
+ Target mainOutputTarget =
+ Iterables.getOnlyElement(
+ Sets.difference(descriptor.getOutputTargetCoders().keySet(), timerOutputTargets));
+ assertThat(
+ outputValues.get(mainOutputTarget),
+ containsInAnyOrder(
+ WindowedValue.valueInGlobalWindow(kvBytes("mainX", "")),
+ WindowedValue.valueInGlobalWindow(kvBytes("event", "")),
+ WindowedValue.valueInGlobalWindow(kvBytes("processing", ""))));
+ assertThat(
+ timerStructuralValues(outputValues.get(eventTimeOutputTarget)),
+ containsInAnyOrder(
+ timerStructuralValue(WindowedValue.valueInGlobalWindow(timerBytes("X", 1L))),
+ timerStructuralValue(WindowedValue.valueInGlobalWindow(timerBytes("Y", 11L))),
+ timerStructuralValue(WindowedValue.valueInGlobalWindow(timerBytes("Z", 21L)))));
+ assertThat(
+ timerStructuralValues(outputValues.get(processingTimeOutputTarget)),
+ containsInAnyOrder(
+ timerStructuralValue(WindowedValue.valueInGlobalWindow(timerBytes("X", 2L))),
+ timerStructuralValue(WindowedValue.valueInGlobalWindow(timerBytes("Y", 12L))),
+ timerStructuralValue(WindowedValue.valueInGlobalWindow(timerBytes("Z", 22L)))));
+ }
+
private KV<byte[], byte[]> kvBytes(String key, long value) throws CoderException {
return KV.of(
CoderUtils.encodeToByteArray(StringUtf8Coder.of(), key),
@@ -553,4 +720,26 @@ public class RemoteExecutionTest implements Serializable {
CoderUtils.encodeToByteArray(StringUtf8Coder.of(), key),
CoderUtils.encodeToByteArray(StringUtf8Coder.of(), value));
}
+
+ private KV<byte[], org.apache.beam.runners.core.construction.Timer<byte[]>> timerBytes(
+ String key, long timestampOffset) throws CoderException {
+ return KV.of(
+ CoderUtils.encodeToByteArray(StringUtf8Coder.of(), key),
+ org.apache.beam.runners.core.construction.Timer.of(
+ BoundedWindow.TIMESTAMP_MIN_VALUE.plus(timestampOffset),
+ CoderUtils.encodeToByteArray(VoidCoder.of(), null, Coder.Context.NESTED)));
+ }
+
+ private Object timerStructuralValue(Object timer) {
+ return WindowedValue.FullWindowedValueCoder.of(
+ KvCoder.of(
+ ByteArrayCoder.of(),
+ org.apache.beam.runners.core.construction.Timer.Coder.of(ByteArrayCoder.of())),
+ GlobalWindow.Coder.INSTANCE)
+ .structuralValue(timer);
+ }
+
+ private Collection<Object> timerStructuralValues(Collection<?> timers) {
+ return Collections2.transform(timers, this::timerStructuralValue);
+ }
}
diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java
index 1c20d74..cd05930 100644
--- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java
+++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java
@@ -33,10 +33,11 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
import com.google.common.collect.ImmutableMap;
-import java.io.IOException;
+import com.google.common.collect.Iterables;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
+import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
@@ -44,8 +45,6 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi.InstructionResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleDescriptor;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.ProcessBundleResponse;
import org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort;
-import org.apache.beam.model.fnexecution.v1.BeamFnApi.Target;
-import org.apache.beam.model.pipeline.v1.Endpoints;
import org.apache.beam.model.pipeline.v1.Endpoints.ApiServiceDescriptor;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
@@ -97,11 +96,73 @@ public class SdkHarnessClientTest {
@Rule public ExpectedException thrown = ExpectedException.none();
private SdkHarnessClient sdkHarnessClient;
+ private ProcessBundleDescriptor descriptor;
+ private String inputPCollection;
+ private BeamFnApi.Target sdkGrpcReadTarget;
+ private BeamFnApi.Target sdkGrpcWriteTarget;
@Before
- public void setup() {
+ public void setup() throws Exception {
MockitoAnnotations.initMocks(this);
sdkHarnessClient = SdkHarnessClient.usingFnApiClient(fnApiControlClient, dataService);
+
+ Pipeline userPipeline = Pipeline.create();
+ TupleTag<String> outputTag = new TupleTag<>();
+ userPipeline
+ .apply("create", Create.of("foo"))
+ .apply("proc", ParDo.of(new TestFn()).withOutputTags(outputTag, TupleTagList.empty()));
+ RunnerApi.Pipeline userProto = PipelineTranslation.toProto(userPipeline);
+
+ ProcessBundleDescriptor.Builder pbdBuilder =
+ ProcessBundleDescriptor.newBuilder()
+ .setId("my_id")
+ .putAllEnvironments(userProto.getComponents().getEnvironmentsMap())
+ .putAllWindowingStrategies(userProto.getComponents().getWindowingStrategiesMap())
+ .putAllCoders(userProto.getComponents().getCodersMap());
+ RunnerApi.Coder fullValueCoder =
+ CoderTranslation.toProto(WindowedValue.getFullCoder(StringUtf8Coder.of(), Coder.INSTANCE))
+ .getCoder();
+ pbdBuilder.putCoders("wire_coder", fullValueCoder);
+
+ PTransform targetProcessor = userProto.getComponents().getTransformsOrThrow("proc");
+ RemoteGrpcPort port =
+ RemoteGrpcPort.newBuilder()
+ .setApiServiceDescriptor(harness.dataEndpoint())
+ .setCoderId("wire_coder")
+ .build();
+ RemoteGrpcPortRead readNode =
+ RemoteGrpcPortRead.readFromPort(
+ port, getOnlyElement(targetProcessor.getInputsMap().values()));
+ RemoteGrpcPortWrite writeNode =
+ RemoteGrpcPortWrite.writeToPort(
+ getOnlyElement(targetProcessor.getOutputsMap().values()), port);
+ // TODO: Ensure cross-env (Runner <-> SDK GRPC Read/Write Node) coders are length-prefixed
+ for (String pc : targetProcessor.getInputsMap().values()) {
+ pbdBuilder.putPcollections(pc, userProto.getComponents().getPcollectionsOrThrow(pc));
+ }
+ for (String pc : targetProcessor.getOutputsMap().values()) {
+ pbdBuilder.putPcollections(pc, userProto.getComponents().getPcollectionsOrThrow(pc));
+ }
+ pbdBuilder
+ .putTransforms("proc", targetProcessor)
+ .putTransforms("read", readNode.toPTransform())
+ .putTransforms("write", writeNode.toPTransform());
+ descriptor = pbdBuilder.build();
+
+ inputPCollection =
+ getOnlyElement(descriptor.getTransformsOrThrow("read").getOutputsMap().values());
+ sdkGrpcReadTarget =
+ BeamFnApi.Target.newBuilder()
+ .setName(
+ getOnlyElement(descriptor.getTransformsOrThrow("read").getOutputsMap().keySet()))
+ .setPrimitiveTransformReference("read")
+ .build();
+ sdkGrpcWriteTarget =
+ BeamFnApi.Target.newBuilder()
+ .setName(
+ getOnlyElement(descriptor.getTransformsOrThrow("write").getInputsMap().keySet()))
+ .setPrimitiveTransformReference("write")
+ .build();
}
@Test
@@ -115,13 +176,16 @@ public class SdkHarnessClientTest {
ProcessBundleDescriptor descriptor2 =
ProcessBundleDescriptor.newBuilder().setId("descriptor2").build();
- RemoteInputDestination<WindowedValue<Integer>> remoteInputs =
- RemoteInputDestination.of(
- FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE),
- Target.getDefaultInstance());
+ Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputs =
+ Collections.singletonMap(
+ "inputPC",
+ RemoteInputDestination.of(
+ (FullWindowedValueCoder)
+ FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE),
+ sdkGrpcReadTarget));
- BundleProcessor<?> processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs);
- BundleProcessor<?> processor2 = sdkHarnessClient.getProcessor(descriptor2, remoteInputs);
+ BundleProcessor processor1 = sdkHarnessClient.getProcessor(descriptor1, remoteInputs);
+ BundleProcessor processor2 = sdkHarnessClient.getProcessor(descriptor2, remoteInputs);
assertNotSame(processor1, processor2);
@@ -141,10 +205,13 @@ public class SdkHarnessClientTest {
.setStateApiServiceDescriptor(ApiServiceDescriptor.newBuilder().setUrl("foo"))
.build();
- RemoteInputDestination<WindowedValue<Integer>> remoteInputs =
- RemoteInputDestination.of(
- FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE),
- Target.getDefaultInstance());
+ Map<String, RemoteInputDestination<WindowedValue<?>>> remoteInputs =
+ Collections.singletonMap(
+ "inputPC",
+ RemoteInputDestination.of(
+ (FullWindowedValueCoder)
+ FullWindowedValueCoder.of(VarIntCoder.of(), GlobalWindow.Coder.INSTANCE),
+ sdkGrpcReadTarget));
thrown.expect(IllegalStateException.class);
thrown.expectMessage("containing a state");
@@ -153,10 +220,6 @@ public class SdkHarnessClientTest {
@Test
public void testNewBundleNoDataDoesNotCrash() throws Exception {
- String descriptorId1 = "descriptor1";
-
- ProcessBundleDescriptor descriptor =
- ProcessBundleDescriptor.newBuilder().setId(descriptorId1).build();
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(new CompletableFuture<>())
@@ -164,12 +227,15 @@ public class SdkHarnessClientTest {
FullWindowedValueCoder<String> coder =
FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
- BundleProcessor<String> processor =
+ BundleProcessor processor =
sdkHarnessClient.getProcessor(
- descriptor, RemoteInputDestination.of(coder, Target.getDefaultInstance()));
+ descriptor,
+ Collections.singletonMap(
+ "inputPC",
+ RemoteInputDestination.of((FullWindowedValueCoder) coder, sdkGrpcReadTarget)));
when(dataService.send(any(), eq(coder))).thenReturn(mock(CloseableFnDataReceiver.class));
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(Collections.emptyMap(), BundleProgressHandler.unsupported())) {
// Correlating the ProcessBundleRequest and ProcessBundleResponse is owned by the underlying
// FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
@@ -185,34 +251,19 @@ public class SdkHarnessClientTest {
@Test
public void testNewBundleAndProcessElements() throws Exception {
- ProcessBundleDescriptor processBundleDescriptor =
- getProcessBundleDescriptor(harness.dataEndpoint());
-
- BeamFnApi.Target sdkGrpcReadTarget =
- BeamFnApi.Target.newBuilder()
- .setName(
- getOnlyElement(
- processBundleDescriptor.getTransformsOrThrow("read").getOutputsMap().keySet()))
- .setPrimitiveTransformReference("read")
- .build();
- BeamFnApi.Target sdkGrpcWriteTarget =
- BeamFnApi.Target.newBuilder()
- .setName(
- getOnlyElement(
- processBundleDescriptor.getTransformsOrThrow("write").getInputsMap().keySet()))
- .setPrimitiveTransformReference("write")
- .build();
-
SdkHarnessClient client = harness.client();
- BundleProcessor<String> processor =
+ BundleProcessor processor =
client.getProcessor(
- processBundleDescriptor,
- RemoteInputDestination.of(
- FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE),
- sdkGrpcReadTarget));
+ descriptor,
+ Collections.singletonMap(
+ "inputPC",
+ RemoteInputDestination.of(
+ (FullWindowedValueCoder)
+ FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE),
+ sdkGrpcReadTarget)));
Collection<WindowedValue<String>> outputs = new ArrayList<>();
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(
Collections.singletonMap(
sdkGrpcWriteTarget,
@@ -221,7 +272,8 @@ public class SdkHarnessClientTest {
LengthPrefixCoder.of(StringUtf8Coder.of()), Coder.INSTANCE),
outputs::add)),
BundleProgressHandler.unsupported())) {
- FnDataReceiver<WindowedValue<String>> bundleInputReceiver = activeBundle.getInputReceiver();
+ FnDataReceiver<WindowedValue<?>> bundleInputReceiver =
+ Iterables.getOnlyElement(activeBundle.getInputReceivers().values());
bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("foo"));
bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("bar"));
bundleInputReceiver.accept(WindowedValue.valueInGlobalWindow("baz"));
@@ -238,14 +290,11 @@ public class SdkHarnessClientTest {
@Test
public void handleCleanupWhenInputSenderFails() throws Exception {
- String descriptorId1 = "descriptor1";
Exception testException = new Exception();
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
- ProcessBundleDescriptor descriptor =
- ProcessBundleDescriptor.newBuilder().setId(descriptorId1).build();
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(new CompletableFuture<>())
@@ -253,9 +302,12 @@ public class SdkHarnessClientTest {
FullWindowedValueCoder<String> coder =
FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
- BundleProcessor<String> processor =
+ BundleProcessor processor =
sdkHarnessClient.getProcessor(
- descriptor, RemoteInputDestination.of(coder, Target.getDefaultInstance()));
+ descriptor,
+ Collections.singletonMap(
+ "inputPC",
+ RemoteInputDestination.of((FullWindowedValueCoder) coder, sdkGrpcReadTarget)));
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
@@ -265,10 +317,9 @@ public class SdkHarnessClientTest {
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
try {
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(
- ImmutableMap.of(Target.getDefaultInstance(), mockRemoteOutputReceiver),
- mockProgressHandler)) {
+ ImmutableMap.of(sdkGrpcWriteTarget, mockRemoteOutputReceiver), mockProgressHandler)) {
// We shouldn't be required to complete the process bundle response future.
}
fail("Exception expected");
@@ -282,7 +333,6 @@ public class SdkHarnessClientTest {
@Test
public void handleCleanupWithStateWhenInputSenderFails() throws Exception {
- String descriptorId1 = "descriptor1";
Exception testException = new Exception();
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
@@ -295,8 +345,6 @@ public class SdkHarnessClientTest {
StateRequestHandler mockStateHandler = mock(StateRequestHandler.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
- ProcessBundleDescriptor descriptor =
- ProcessBundleDescriptor.newBuilder().setId(descriptorId1).build();
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(new CompletableFuture<>())
@@ -304,10 +352,12 @@ public class SdkHarnessClientTest {
FullWindowedValueCoder<String> coder =
FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
- BundleProcessor<String> processor =
+ BundleProcessor processor =
sdkHarnessClient.getProcessor(
descriptor,
- RemoteInputDestination.of(coder, Target.getDefaultInstance()),
+ Collections.singletonMap(
+ inputPCollection,
+ RemoteInputDestination.of((FullWindowedValueCoder) coder, sdkGrpcReadTarget)),
mockStateDelegator);
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
@@ -317,9 +367,9 @@ public class SdkHarnessClientTest {
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
try {
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(
- ImmutableMap.of(Target.getDefaultInstance(), mockRemoteOutputReceiver),
+ ImmutableMap.of(sdkGrpcWriteTarget, mockRemoteOutputReceiver),
mockStateHandler,
mockProgressHandler)) {
// We shouldn't be required to complete the process bundle response future.
@@ -341,8 +391,6 @@ public class SdkHarnessClientTest {
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
- ProcessBundleDescriptor descriptor =
- ProcessBundleDescriptor.newBuilder().setId("descriptor1").build();
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(new CompletableFuture<>())
@@ -350,9 +398,12 @@ public class SdkHarnessClientTest {
FullWindowedValueCoder<String> coder =
FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
- BundleProcessor<String> processor =
+ BundleProcessor processor =
sdkHarnessClient.getProcessor(
- descriptor, RemoteInputDestination.of(coder, Target.getDefaultInstance()));
+ descriptor,
+ Collections.singletonMap(
+ "inputPC",
+ RemoteInputDestination.of((FullWindowedValueCoder) coder, sdkGrpcReadTarget)));
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
@@ -360,10 +411,9 @@ public class SdkHarnessClientTest {
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
try {
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(
- ImmutableMap.of(Target.getDefaultInstance(), mockRemoteOutputReceiver),
- mockProgressHandler)) {
+ ImmutableMap.of(sdkGrpcWriteTarget, mockRemoteOutputReceiver), mockProgressHandler)) {
processBundleResponseFuture.completeExceptionally(testException);
}
fail("Exception expected");
@@ -377,7 +427,6 @@ public class SdkHarnessClientTest {
@Test
public void handleCleanupWithStateWhenProcessingBundleFails() throws Exception {
- String descriptorId1 = "descriptor1";
Exception testException = new Exception();
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
@@ -389,8 +438,6 @@ public class SdkHarnessClientTest {
StateRequestHandler mockStateHandler = mock(StateRequestHandler.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
- ProcessBundleDescriptor descriptor =
- ProcessBundleDescriptor.newBuilder().setId(descriptorId1).build();
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(new CompletableFuture<>())
@@ -398,10 +445,12 @@ public class SdkHarnessClientTest {
FullWindowedValueCoder<String> coder =
FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
- BundleProcessor<String> processor =
+ BundleProcessor processor =
sdkHarnessClient.getProcessor(
descriptor,
- RemoteInputDestination.of(coder, Target.getDefaultInstance()),
+ Collections.singletonMap(
+ inputPCollection,
+ RemoteInputDestination.of((FullWindowedValueCoder) coder, sdkGrpcReadTarget)),
mockStateDelegator);
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
@@ -409,9 +458,9 @@ public class SdkHarnessClientTest {
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
try {
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(
- ImmutableMap.of(Target.getDefaultInstance(), mockRemoteOutputReceiver),
+ ImmutableMap.of(sdkGrpcWriteTarget, mockRemoteOutputReceiver),
mockStateHandler,
mockProgressHandler)) {
processBundleResponseFuture.completeExceptionally(testException);
@@ -428,14 +477,11 @@ public class SdkHarnessClientTest {
@Test
public void handleCleanupWhenAwaitingOnClosingOutputReceivers() throws Exception {
- String descriptorId1 = "descriptor1";
Exception testException = new Exception();
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
CloseableFnDataReceiver mockInputSender = mock(CloseableFnDataReceiver.class);
- ProcessBundleDescriptor descriptor =
- ProcessBundleDescriptor.newBuilder().setId(descriptorId1).build();
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(new CompletableFuture<>())
@@ -443,9 +489,12 @@ public class SdkHarnessClientTest {
FullWindowedValueCoder<String> coder =
FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
- BundleProcessor<String> processor =
+ BundleProcessor processor =
sdkHarnessClient.getProcessor(
- descriptor, RemoteInputDestination.of(coder, Target.getDefaultInstance()));
+ descriptor,
+ Collections.singletonMap(
+ "inputPC",
+ RemoteInputDestination.of((FullWindowedValueCoder) coder, sdkGrpcReadTarget)));
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
doThrow(testException).when(mockOutputReceiver).awaitCompletion();
@@ -454,10 +503,9 @@ public class SdkHarnessClientTest {
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
try {
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(
- ImmutableMap.of(Target.getDefaultInstance(), mockRemoteOutputReceiver),
- mockProgressHandler)) {
+ ImmutableMap.of(sdkGrpcWriteTarget, mockRemoteOutputReceiver), mockProgressHandler)) {
// Correlating the ProcessBundleRequest and ProcessBundleResponse is owned by the underlying
// FnApiControlClient. The SdkHarnessClient owns just wrapping the request and unwrapping
// the response.
@@ -477,7 +525,6 @@ public class SdkHarnessClientTest {
@Test
public void handleCleanupWithStateWhenAwaitingOnClosingOutputReceivers() throws Exception {
- String descriptorId1 = "descriptor1";
Exception testException = new Exception();
InboundDataClient mockOutputReceiver = mock(InboundDataClient.class);
@@ -489,8 +536,6 @@ public class SdkHarnessClientTest {
StateRequestHandler mockStateHandler = mock(StateRequestHandler.class);
BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class);
- ProcessBundleDescriptor descriptor =
- ProcessBundleDescriptor.newBuilder().setId(descriptorId1).build();
CompletableFuture<InstructionResponse> processBundleResponseFuture = new CompletableFuture<>();
when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class)))
.thenReturn(new CompletableFuture<>())
@@ -498,10 +543,12 @@ public class SdkHarnessClientTest {
FullWindowedValueCoder<String> coder =
FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE);
- BundleProcessor<String> processor =
+ BundleProcessor processor =
sdkHarnessClient.getProcessor(
descriptor,
- RemoteInputDestination.of(coder, Target.getDefaultInstance()),
+ Collections.singletonMap(
+ inputPCollection,
+ RemoteInputDestination.of((FullWindowedValueCoder) coder, sdkGrpcReadTarget)),
mockStateDelegator);
when(dataService.receive(any(), any(), any())).thenReturn(mockOutputReceiver);
when(dataService.send(any(), eq(coder))).thenReturn(mockInputSender);
@@ -510,9 +557,9 @@ public class SdkHarnessClientTest {
RemoteOutputReceiver mockRemoteOutputReceiver = mock(RemoteOutputReceiver.class);
try {
- try (ActiveBundle<String> activeBundle =
+ try (ActiveBundle activeBundle =
processor.newBundle(
- ImmutableMap.of(Target.getDefaultInstance(), mockRemoteOutputReceiver),
+ ImmutableMap.of(sdkGrpcWriteTarget, mockRemoteOutputReceiver),
mockStateHandler,
mockProgressHandler)) {
// Correlating the ProcessBundleRequest and ProcessBundleResponse is owned by the underlying
@@ -532,52 +579,6 @@ public class SdkHarnessClientTest {
}
}
- private BeamFnApi.ProcessBundleDescriptor getProcessBundleDescriptor(
- Endpoints.ApiServiceDescriptor endpoint) throws IOException {
- Pipeline userPipeline = Pipeline.create();
- TupleTag<String> outputTag = new TupleTag<>();
- userPipeline
- .apply("create", Create.of("foo"))
- .apply("proc", ParDo.of(new TestFn()).withOutputTags(outputTag, TupleTagList.empty()));
- RunnerApi.Pipeline userProto = PipelineTranslation.toProto(userPipeline);
-
- ProcessBundleDescriptor.Builder pbdBuilder =
- ProcessBundleDescriptor.newBuilder()
- .setId("my_id")
- .putAllEnvironments(userProto.getComponents().getEnvironmentsMap())
- .putAllWindowingStrategies(userProto.getComponents().getWindowingStrategiesMap())
- .putAllCoders(userProto.getComponents().getCodersMap());
- RunnerApi.Coder fullValueCoder =
- CoderTranslation.toProto(WindowedValue.getFullCoder(StringUtf8Coder.of(), Coder.INSTANCE))
- .getCoder();
- pbdBuilder.putCoders("wire_coder", fullValueCoder);
-
- PTransform targetProcessor = userProto.getComponents().getTransformsOrThrow("proc");
- RemoteGrpcPort port =
- RemoteGrpcPort.newBuilder()
- .setApiServiceDescriptor(endpoint)
- .setCoderId("wire_coder")
- .build();
- RemoteGrpcPortRead readNode =
- RemoteGrpcPortRead.readFromPort(
- port, getOnlyElement(targetProcessor.getInputsMap().values()));
- RemoteGrpcPortWrite writeNode =
- RemoteGrpcPortWrite.writeToPort(
- getOnlyElement(targetProcessor.getOutputsMap().values()), port);
- // TODO: Ensure cross-env (Runner <-> SDK GRPC Read/Write Node) coders are length-prefixed
- for (String pc : targetProcessor.getInputsMap().values()) {
- pbdBuilder.putPcollections(pc, userProto.getComponents().getPcollectionsOrThrow(pc));
- }
- for (String pc : targetProcessor.getOutputsMap().values()) {
- pbdBuilder.putPcollections(pc, userProto.getComponents().getPcollectionsOrThrow(pc));
- }
- pbdBuilder
- .putTransforms("proc", targetProcessor)
- .putTransforms("read", readNode.toPTransform())
- .putTransforms("write", writeNode.toPTransform());
- return pbdBuilder.build();
- }
-
private static class TestFn extends DoFn<String, String> {
@ProcessElement
public void processElement(ProcessContext context) {