/** * Update all composites present in the {@code originalPipeline} with an URN equal to the provided * {@code urn} using the provided {@link TransformReplacement}. */ public static Pipeline updateTransform( String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder) { Components.Builder resultComponents = originalPipeline.getComponents().toBuilder(); for (Map.Entry<String, PTransform> pt : originalPipeline.getComponents().getTransformsMap().entrySet()) { if (pt.getValue().getSpec() != null && urn.equals(pt.getValue().getSpec().getUrn())) { MessageWithComponents updated = compositeBuilder.getReplacement(pt.getKey(), originalPipeline.getComponents()); checkArgument( updated.getPtransform().getOutputsMap().equals(pt.getValue().getOutputsMap()), "A %s must produce all of the outputs of the original %s", TransformReplacement.class.getSimpleName(), PTransform.class.getSimpleName()); removeSubtransforms(pt.getValue(), resultComponents); resultComponents .mergeFrom(updated.getComponents()) .putTransforms(pt.getKey(), updated.getPtransform()); } } return originalPipeline.toBuilder().setComponents(resultComponents).build(); }
public ExecutableGraphBuilder addTransform( String name, @Nullable String input, String... outputs) { PTransform.Builder pt = PTransform.newBuilder().setUniqueName(name); if (input != null) { pt = pt.putInputs("input", input); addPCollection(input); } for (String output : outputs) { pt = pt.putOutputs(output, output); addPCollection(output); } components.putTransforms(name, pt.build()); return this; }
@Test public void forTransformsWithSubgraph() { Components components = Components.newBuilder() .putTransforms( "root", PTransform.newBuilder().putOutputs("output", "output.out").build()) .putPcollections( "output.out", RunnerApi.PCollection.newBuilder().setUniqueName("output.out").build()) .putTransforms( "consumer", PTransform.newBuilder().putInputs("input", "output.out").build()) .putTransforms( "ignored", PTransform.newBuilder().putInputs("input", "output.out").build()) .build(); QueryablePipeline pipeline = QueryablePipeline.forTransforms(ImmutableSet.of("root", "consumer"), components); assertThat( pipeline.getRootTransforms(), contains(PipelineNode.pTransform("root", components.getTransformsOrThrow("root")))); Set<PTransformNode> consumers = pipeline.getPerElementConsumers( PipelineNode.pCollection( "output.out", components.getPcollectionsOrThrow("output.out"))); assertThat( consumers, contains(PipelineNode.pTransform("consumer", components.getTransformsOrThrow("consumer")))); }
partialComponents .toBuilder() .putTransforms( "read", PTransform.newBuilder() .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( "goTransform", PTransform.newBuilder() .build()) .putPcollections("go.out", PCollection.newBuilder().setUniqueName("go.out").build()) .putTransforms( "pyTransform", PTransform.newBuilder()
partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections("read.out", userStateMainInputPCollection) .putTransforms( "user_state", PTransform.newBuilder() "user_state.out", PCollection.newBuilder().setUniqueName("user_state.out").build()) .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build())
partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections("parDo.out", pc("parDo.out")) .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out")) .putPcollections("output.out", pc("output.out"))
partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms("timer", timerTransform) .putPcollections( "timer.out", PCollection.newBuilder().setUniqueName("timer.out").build())
.putPcollections("read_pc", RunnerApi.PCollection.getDefaultInstance()) .putPcollections("pardo_out", RunnerApi.PCollection.getDefaultInstance()) .putTransforms( "root", PTransform.newBuilder() .putOutputs("out", "read_pc") .build()) .putTransforms( "multiConsumer", PTransform.newBuilder()
partialComponents .toBuilder() .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms("window", windowTransform) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build())
.setComponents( Components.newBuilder() .putTransforms( "root", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:composite")) .build()) .putTransforms( "sub_first", PTransform.newBuilder()
partialComponents .toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( "gbk", PTransform.newBuilder()
partialComponents .toBuilder() .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out")) .putPcollections("output.out", pc("output.out"))
partialComponents .toBuilder() .putTransforms( "mystery", PTransform.newBuilder() .build()) .putPcollections("mystery.out", pc("mystery.out")) .putTransforms( "enigma", PTransform.newBuilder()
@Test public void noEnvironmentThrows() { // (impulse.out) -> runnerTransform -> gbk.out // runnerTransform can't be executed in an environment, so trying to construct it should fail PTransform gbkTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .setSpec( FunctionSpec.newBuilder().setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)) .putOutputs("output", "gbk.out") .build(); QueryablePipeline p = QueryablePipeline.forPrimitivesIn( partialComponents .toBuilder() .putTransforms("runnerTransform", gbkTransform) .putPcollections( "gbk.out", PCollection.newBuilder().setUniqueName("gbk.out").build()) .build()); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("Environment must be populated"); GreedyStageFuser.forGrpcPortRead( p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("runnerTransform", gbkTransform))); }
@Before public void setup() { partialComponents = Components.newBuilder() .putTransforms( "impulse", PTransform.newBuilder() .setUniqueName("Impulse") .putOutputs("output", "impulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("impulse.out", pc("impulse.out")) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build()) .build(); }
/** * Registers the provided {@link AppliedPTransform} into this {@link SdkComponents}, returning a * unique ID for the {@link AppliedPTransform}. Multiple registrations of the same {@link * AppliedPTransform} will return the same unique ID. * * <p>All of the children must already be registered within this {@link SdkComponents}. */ public String registerPTransform( AppliedPTransform<?, ?, ?> appliedPTransform, List<AppliedPTransform<?, ?, ?>> children) throws IOException { String name = getApplicationName(appliedPTransform); // If this transform is present in the components, nothing to do. return the existing name. // Otherwise the transform must be translated and added to the components. if (componentsBuilder.getTransformsOrDefault(name, null) != null) { return name; } checkNotNull(children, "child nodes may not be null"); componentsBuilder.putTransforms( name, PTransformTranslation.toProto(appliedPTransform, children, this)); return name; }
@Test public void fromComponentsWithMalformedComponents() { Components components = Components.newBuilder() .putTransforms( "root", PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN) .build()) .putOutputs("output", "output.out") .build()) .build(); thrown.expect(IllegalArgumentException.class); QueryablePipeline.forPrimitivesIn(components).getComponents(); }
@Before public void setup() { partialComponents = Components.newBuilder() .putTransforms( "impulse", PTransform.newBuilder() .putOutputs("output", "impulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("impulse.out", impulseDotOut) .build(); }
@Test public void forTransformsWithMalformedGraph() { Components components = Components.newBuilder() .putTransforms( "root", PTransform.newBuilder().putOutputs("output", "output.out").build()) .putPcollections( "output.out", RunnerApi.PCollection.newBuilder().setUniqueName("output.out").build()) .putTransforms( "consumer", PTransform.newBuilder().putInputs("input", "output.out").build()) .build(); thrown.expect(IllegalArgumentException.class); // Consumer consumes a PCollection which isn't produced. QueryablePipeline.forTransforms(ImmutableSet.of("consumer"), components); }
private void removeDescendants(RunnerApi.Pipeline.Builder pipeline, String parentId) { RunnerApi.PTransform parentProto = pipeline.getComponents().getTransformsOrDefault(parentId, null); if (parentProto != null) { for (String childId : parentProto.getSubtransformsList()) { removeDescendants(pipeline, childId); pipeline.getComponentsBuilder().removeTransforms(childId); } pipeline .getComponentsBuilder() .putTransforms(parentId, parentProto.toBuilder().clearSubtransforms().build()); } }