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(); }
@Test public void getEnvironmentUnknownFnType() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN) .build()) .build(); Optional<Environment> env = Environments.getEnvironment(builder, rehydratedComponents); assertThat(env.isPresent(), is(false)); }
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; }
PTransform.newBuilder() .putInputs("input", "impulse.out") .putOutputs("output", "read.out") .putTransforms( "gbk", PTransform.newBuilder() .putInputs("input", "read.out") .putOutputs("output", "gbk.out")
RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder()
@Test public void parDoWithStateAndTimerRootsStage() { PTransform timerTransform = PTransform.newBuilder() .setUniqueName("TimerParDo") .putInputs("input", "impulse.out")
.putTransforms( "mystery", PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .putTransforms( "enigma", PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder()
@Test public void getEnvironmentWindowIntoKnown() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowIntoPayload payload = WindowIntoPayload.newBuilder() .setWindowFn( WindowingStrategyTranslation.toProto( FixedWindows.of(Duration.standardMinutes(5L)), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getWindowFn().getEnvironmentId()))); }
RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder()
@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 getEnvironmentCombine() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); CombinePayload payload = CombinePayload.newBuilder() .setCombineFn(CombineTranslation.toProto(Sum.ofLongs(), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getCombineFn().getEnvironmentId()))); } }
@Test public void getEnvironmentRead() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); ReadPayload payload = ReadTranslation.toProto(Read.from(CountingSource.unbounded()), components); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getSource().getEnvironmentId()))); }
@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()); }
@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(); }
@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(); }
@Test public void testMatch() { Assert.assertTrue( NativeTransforms.isNative( RunnerApi.PTransform.newBuilder() .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn("test").build()) .build())); }