/** * Returns a {@link PTransform} like the input {@link PTransform}, but with each output to {@code * originalPCollection} replaced with an output (with the same local name) to {@code * newPCollection}. */ private static PTransform updateOutputs( PTransform transform, Map<String, PCollectionNode> originalToPartial) { PTransform.Builder updatedTransformBuilder = transform.toBuilder(); for (Map.Entry<String, String> output : transform.getOutputsMap().entrySet()) { if (originalToPartial.containsKey(output.getValue())) { updatedTransformBuilder.putOutputs( output.getKey(), originalToPartial.get(output.getValue()).getId()); } } return updatedTransformBuilder.build(); }
"Unexpected output type %s", taggedOutput.getValue().getClass()); transformBuilder.putOutputs( toProto(taggedOutput.getKey()), components.registerPCollection((PCollection<?>) taggedOutput.getValue()));
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(); }
.setUniqueName("ParDo") .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() .putInputs("input", "parDo.out") .putInputs("timer", "timer.out") .putOutputs("timer", "timer.out") .putOutputs("output", "output.out") .setSpec( FunctionSpec.newBuilder()
PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() PTransform.newBuilder() .putInputs("input", "parDo.out") .putOutputs("output", "timer.out") .setSpec( FunctionSpec.newBuilder()
.setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN) .build()) .putOutputs("out", "read_pc") .build()) .putTransforms( .putInputs("main_in", "read_pc") .putInputs("side_in", "read_pc") .putOutputs("out", "pardo_out") .setSpec( FunctionSpec.newBuilder()
PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "parDo.out") .setSpec( FunctionSpec.newBuilder() PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "window.out") .setSpec( FunctionSpec.newBuilder()
String timerPCollectionId = components.registerPCollection(timerPCollection); builder.putInputs(localTimerName, timerPCollectionId); builder.putOutputs(localTimerName, timerPCollectionId);
PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "gbk.out") .setSpec( FunctionSpec.newBuilder()
.putInputs("input", "impulse.out") .putInputs("timer", "timer.out") .putOutputs("timer", "timer.out") .putOutputs("output", "output.out") .setSpec( FunctionSpec.newBuilder()
.setUniqueName("Mystery") .putInputs("input", "impulse.out") .putOutputs("output", "mystery.out") .build()) .putPcollections("mystery.out", pc("mystery.out")) .setUniqueName("Enigma") .putInputs("input", "impulse.out") .putOutputs("output", "enigma.out") .build()) .putPcollections("enigma.out", pc("enigma.out"))
@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() { 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(); }
@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() { 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()); }
@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); }
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; }
@Before public void setup() { pc = PipelineNode.pCollection( "pc", RunnerApi.PCollection.newBuilder().setUniqueName("pc").build()); transform = PipelineNode.pTransform("pt", PTransform.newBuilder().putOutputs("out", "pc").build()); bundleFactory = ImmutableListBundleFactory.create(); }