Examples with PTransform org.apache.beam.sdk.common.runner.v1.RunnerApi.PTransform used on opensource projects

Example 56 with PTransform

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.PTransform in project beam by apache.

the class GreedyStageFuserTest method materializesWithStatefulConsumer.

@Test
public void materializesWithStatefulConsumer() {
    // (impulse.out) -> parDo -> (parDo.out)
    // (parDo.out) -> stateful -> stateful.out
    // stateful has a state spec which prevents it from fusing with an upstream ParDo
    PTransform parDoTransform = PTransform.newBuilder().putInputs("input", "impulse.out").putOutputs("output", "parDo.out").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN).setPayload(ParDoPayload.newBuilder().setDoFn(FunctionSpec.newBuilder()).build().toByteString())).setEnvironmentId("common").build();
    PTransform statefulTransform = PTransform.newBuilder().putInputs("input", "parDo.out").putOutputs("output", "stateful.out").setSpec(FunctionSpec.newBuilder().setUrn(PTransformTranslation.PAR_DO_TRANSFORM_URN).setPayload(ParDoPayload.newBuilder().setDoFn(FunctionSpec.newBuilder()).putStateSpecs("state", StateSpec.getDefaultInstance()).build().toByteString())).setEnvironmentId("common").build();
    QueryablePipeline p = QueryablePipeline.forPrimitivesIn(partialComponents.toBuilder().putTransforms("parDo", parDoTransform).putPcollections("parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build()).putTransforms("stateful", statefulTransform).putPcollections("stateful.out", PCollection.newBuilder().setUniqueName("stateful.out").build()).putEnvironments("common", Environments.createDockerEnvironment("common")).build());
    ExecutableStage subgraph = GreedyStageFuser.forGrpcPortRead(p, impulseOutputNode, ImmutableSet.of(PipelineNode.pTransform("parDo", parDoTransform)));
    assertThat(subgraph.getOutputPCollections(), contains(PipelineNode.pCollection("parDo.out", PCollection.newBuilder().setUniqueName("parDo.out").build())));
    assertThat(subgraph, hasSubtransforms("parDo"));
}

Example 57 with PTransform

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.PTransform in project beam by apache.

the class EnvironmentsTest method getEnvironmentUnknownFnType.

@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));
}

Example 58 with PTransform

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.PTransform in project beam by apache.

the class ProtoOverrides method removeSubtransforms.

/**
 * Remove all sub-transforms of the provided transform recursively. A {@link PTransform} can be
 * the sub-transform of only one enclosing transform.
 */
private static void removeSubtransforms(PTransform pt, Components.Builder target) {
    for (String subtransformId : pt.getSubtransformsList()) {
        PTransform subtransform = target.getTransformsOrThrow(subtransformId);
        removeSubtransforms(subtransform, target);
        target.removeTransforms(subtransformId);
    // TODO: remove PCollections not produced by 'pt' here.
    }
}

Example 59 with PTransform

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.PTransform in project beam by apache.

the class ProtoOverrides method updateTransform.

/**
 * Update all composites present in the {@code originalPipeline} with an URN equal to the provided
 * {@code urn} using the provided {@link TransformReplacement}.
 */
public static Pipeline updateTransform(String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder) {
    Components.Builder resultComponents = originalPipeline.getComponents().toBuilder();
    for (Map.Entry<String, PTransform> pt : originalPipeline.getComponents().getTransformsMap().entrySet()) {
        if (pt.getValue().getSpec() != null && urn.equals(pt.getValue().getSpec().getUrn())) {
            MessageWithComponents updated = compositeBuilder.getReplacement(pt.getKey(), originalPipeline.getComponents());
            if (updated == null) {
                continue;
            }
            checkArgument(updated.getPtransform().getOutputsMap().equals(pt.getValue().getOutputsMap()), "A %s must produce all of the outputs of the original %s", TransformReplacement.class.getSimpleName(), PTransform.class.getSimpleName());
            removeSubtransforms(pt.getValue(), resultComponents);
            resultComponents.mergeFrom(updated.getComponents()).putTransforms(pt.getKey(), updated.getPtransform());
        }
    }
    return originalPipeline.toBuilder().setComponents(resultComponents).build();
}

Example 60 with PTransform

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.PTransform in project beam by apache.

the class UserStateReference method fromUserStateId.

/**
 * Create a user state reference from a UserStateId proto and components.
 */
public static UserStateReference fromUserStateId(UserStateId userStateId, RunnerApi.Components components) {
    PTransform transform = components.getTransformsOrThrow(userStateId.getTransformId());
    String mainInputCollectionId;
    try {
        mainInputCollectionId = transform.getInputsOrThrow(ParDoTranslation.getMainInputName(transform));
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
    return UserStateReference.of(PipelineNode.pTransform(userStateId.getTransformId(), transform), userStateId.getLocalName(), PipelineNode.pCollection(mainInputCollectionId, components.getPcollectionsOrThrow(mainInputCollectionId)));
}