/** * 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(); }
@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(); }
.setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat(
@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(); }
@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(); } }
@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)); }
@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 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(); }
@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(); }
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; }
private void removeDescendants(RunnerApi.Pipeline.Builder pipeline, String parentId) { RunnerApi.PTransform parentProto = pipeline.getComponents().getTransformsOrDefault(parentId, null); if (parentProto != null) { for (String childId : parentProto.getSubtransformsList()) { removeDescendants(pipeline, childId); pipeline.getComponentsBuilder().removeTransforms(childId); } pipeline .getComponentsBuilder() .putTransforms(parentId, parentProto.toBuilder().clearSubtransforms().build()); } }
@Override public RunnerApi.PTransform translate( AppliedPTransform<?, ?, ?> appliedPTransform, List<AppliedPTransform<?, ?, ?>> subtransforms, SdkComponents components) throws IOException { return translateAppliedPTransform(appliedPTransform, subtransforms, components).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())); }