"Unexpected input type %s", taggedInput.getValue().getClass()); transformBuilder.putInputs( toProto(taggedInput.getKey()), components.registerPCollection((PCollection<?>) taggedInput.getValue()));
.setUrn(SplittableRemoteStageEvaluatorFactory.URN) .build()) .putInputs(
private static PTransform createFlattenOfPartials( String transformId, String outputId, Collection<PCollectionNode> generatedInputs) { PTransform.Builder newFlattenBuilder = PTransform.newBuilder(); int i = 0; for (PCollectionNode generatedInput : generatedInputs) { String localInputId = String.format("input_%s", i); i++; newFlattenBuilder.putInputs(localInputId, generatedInput.getId()); } // Flatten all of the new partial nodes together. return newFlattenBuilder // Use transform ID as unique name. .setUniqueName(transformId) .putOutputs("output", outputId) .setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.FLATTEN_TRANSFORM_URN)) .build(); }
PTransform.newBuilder() .setUniqueName(String.format("%s/GABW", gbk.getUniqueName())) .putInputs("input", kwiCollectionId) .setSpec(FunctionSpec.newBuilder().setUrn(DirectGroupByKey.DIRECT_GABW_URN)) .putAllOutputs(gbk.getOutputsMap())
PTransform readTransform = PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( PTransform parDoTransform = PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "parDo.out") .setSpec( "user_state", PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "user_state.out") .build())
PTransform.newBuilder() .setUniqueName("ParDo") .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( PTransform.newBuilder() .setUniqueName("TimerParDo") .putInputs("input", "parDo.out") .putInputs("timer", "timer.out") .putOutputs("timer", "timer.out") .putOutputs("output", "output.out")
.setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .build()) .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .build()) "goTransform", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "go.out") .setSpec( "pyTransform", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "py.out") .setSpec(
String.format("%s.%s", appliedPTransform.getFullName(), localTimerName)); String timerPCollectionId = components.registerPCollection(timerPCollection); builder.putInputs(localTimerName, timerPCollectionId); builder.putOutputs(localTimerName, timerPCollectionId);
.setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)) .setUniqueName("Mystery") .putInputs("input", "impulse.out") .putOutputs("output", "mystery.out") .build()) .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)) .setUniqueName("Enigma") .putInputs("input", "impulse.out") .putOutputs("output", "enigma.out") .build())
@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")))); }
@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() { create = PipelineNode.pTransform( "create", PTransform.newBuilder() .setUniqueName("create") .putInputs("in", "impulse.out") .putOutputs("out", "create.out") .build()); sum = PipelineNode.pTransform( "sum", PTransform.newBuilder() .setUniqueName("sum") .putInputs("in", "create.in") .putOutputs("out", "sum.out") .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); }
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; }