newComponents .putCoders(kwiCoderId, kwiCoder) .putPcollections(kwiCollectionId, kwi) .putTransforms(gbkoId, gbko);
.putTransforms("pt", pt) .putTransforms("other_pt", PTransform.newBuilder().setUniqueName("other").build()) .putPcollections("input.out", input) .putPcollections("sideInput.in", sideInput) .putPcollections("timer.pc", timer) .putPcollections("output.out", output) .putEnvironments("foo", env) .build();
.setSpec(FunctionSpec.newBuilder().setUrn("beam:repeated")) .build()) .putPcollections( "intermediatePc", PCollection.newBuilder().setUniqueName("intermediate").build())
.toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( .toByteString())) .build()) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putTransforms( .toByteString())) .build()) .putPcollections( "window.out", PCollection.newBuilder().setUniqueName("window.out").build()) .putEnvironments("rare", Environments.createDockerEnvironment("rare"))
.toBuilder() .putTransforms("read", readTransform) .putPcollections("read.out", userStateMainInputPCollection) .putTransforms( "user_state", .putOutputs("output", "user_state.out") .build()) .putPcollections( "user_state.out", PCollection.newBuilder().setUniqueName("user_state.out").build()) .putTransforms("parDo", parDoTransform) .putPcollections( "parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putEnvironments("common", env)
.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()) .putEnvironments("common", Environments.createDockerEnvironment("common"))
.putOutputs("output", "read.out") .build()) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( .toByteString())) .build()) .putPcollections("go.out", PCollection.newBuilder().setUniqueName("go.out").build()) .putTransforms( "pyTransform", .toByteString())) .build()) .putPcollections("py.out", PCollection.newBuilder().setUniqueName("py.out").build()) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py"))
.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()) .putEnvironments("common", Environments.createDockerEnvironment("common"))
Components components = Components.newBuilder() .putPcollections("read_pc", RunnerApi.PCollection.getDefaultInstance()) .putPcollections("pardo_out", RunnerApi.PCollection.getDefaultInstance()) .putTransforms( "root",
.toBuilder() .putTransforms("read", readTransform) .putPcollections( "read.out", PCollection.newBuilder().setUniqueName("read.out").build()) .putTransforms( .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN)) .build()) .putPcollections( "gbk.out", PCollection.newBuilder().setUniqueName("parDo.out").build()) .putEnvironments("common", env)
.toBuilder() .putTransforms("timer", timerTransform) .putPcollections("timer.out", pc("timer.out")) .putPcollections("output.out", pc("output.out")) .putEnvironments("common", Environments.createDockerEnvironment("common")) .build();
.putOutputs("output", "mystery.out") .build()) .putPcollections("mystery.out", pc("mystery.out")) .putTransforms( "enigma", .putOutputs("output", "enigma.out") .build()) .putPcollections("enigma.out", pc("enigma.out")) .build(); FusedPipeline fused =
/** This method doesn't do any pruning for reachability, but this may not require a test. */ @Test public void retainOnlyPrimitivesIgnoresUnreachableNodes() { Pipeline p = Pipeline.create(); p.apply( new org.apache.beam.sdk.transforms.PTransform<PBegin, PCollection<Long>>() { @Override public PCollection<Long> expand(PBegin input) { return input .apply(GenerateSequence.from(2L)) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(5L)))) .apply(MapElements.into(TypeDescriptors.longs()).via(l -> l + 1)); } }); Components augmentedComponents = PipelineTranslation.toProto(p) .getComponents() .toBuilder() .putCoders("extra-coder", RunnerApi.Coder.getDefaultInstance()) .putWindowingStrategies( "extra-windowing-strategy", RunnerApi.WindowingStrategy.getDefaultInstance()) .putEnvironments("extra-env", RunnerApi.Environment.getDefaultInstance()) .putPcollections("extra-pc", RunnerApi.PCollection.getDefaultInstance()) .build(); Collection<String> primitiveComponents = QueryablePipeline.getPrimitiveTransformIds(augmentedComponents); } }
@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(); }
/** * Registers the provided {@link PCollection} into this {@link SdkComponents}, returning a unique * ID for the {@link PCollection}. Multiple registrations of the same {@link PCollection} will * return the same unique ID. */ public String registerPCollection(PCollection<?> pCollection) throws IOException { String existing = pCollectionIds.get(pCollection); if (existing != null) { return existing; } String uniqueName = uniqify(pCollection.getName(), pCollectionIds.values()); pCollectionIds.put(pCollection, uniqueName); componentsBuilder.putPcollections( uniqueName, PCollectionTranslation.toProto(pCollection, this)); return uniqueName; }
@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); }
@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(); }
private ExecutableGraphBuilder addPCollection(String name) { components.putPcollections(name, PCollection.newBuilder().setUniqueName(name).build()); return this; }