@Override public RunnerApi.PTransform translate( AppliedPTransform<?, ?, ?> appliedPTransform, List<AppliedPTransform<?, ?, ?>> subtransforms, SdkComponents components) throws IOException { RunnerApi.PTransform.Builder transformBuilder = translateAppliedPTransform(appliedPTransform, subtransforms, components); FunctionSpec spec = KNOWN_PAYLOAD_TRANSLATORS .get(appliedPTransform.getTransform().getClass()) .translate(appliedPTransform, components); if (spec != null) { transformBuilder.setSpec(spec); } return transformBuilder.build(); } }
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(); }
appliedPTransform.getPipeline(), components); builder.setSpec( RunnerApi.FunctionSpec.newBuilder() .setUrn(PAR_DO_TRANSFORM_URN)
PTransform.newBuilder() .addSubtransforms("sub_first") .setSpec(FunctionSpec.newBuilder().setUrn("beam:composite")) .build()) .putTransforms( "sub_first", PTransform.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn("beam:inner")) .build())) .build(); .addSubtransforms("foo") .addSubtransforms("bar") .setSpec( FunctionSpec.getDefaultInstance() .newBuilderForType()
.putInputs("input", "impulse.out") .putOutputs("output", "read.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN) .putInputs("input", "read.out") .putOutputs("output", "gbk.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN))
.putOutputs("timer", "timer.out") .putOutputs("output", "output.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN)
@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()))); }
PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN)
@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))); }
@Override public RunnerApi.PTransform translate( AppliedPTransform<?, ?, ?> appliedPTransform, List<AppliedPTransform<?, ?, ?>> subtransforms, SdkComponents components) throws IOException { RunnerApi.PTransform.Builder transformBuilder = translateAppliedPTransform(appliedPTransform, subtransforms, components); PTransform<?, ?> transform = appliedPTransform.getTransform(); // The raw transform was parsed in the context of other components; this puts it in the // context of our current serialization FunctionSpec spec = ((RawPTransform<?, ?>) transform).migrate(components); // A composite transform is permitted to have a null spec. There are also some pseudo- // primitives not yet supported by the portability framework that have null specs if (spec != null) { transformBuilder.setSpec(spec); } return transformBuilder.build(); } }
@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 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)); }
@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 testMatch() { Assert.assertTrue( NativeTransforms.isNative( RunnerApi.PTransform.newBuilder() .setSpec(RunnerApi.FunctionSpec.newBuilder().setUrn("test").build()) .build())); }