Examples with Components org.apache.beam.model.pipeline.v1.RunnerApi.Components used on opensource projects

Example 81 with Components

use of org.apache.beam.model.pipeline.v1.RunnerApi.Components in project beam by apache.

the class TestStreamTranslationTest method testRegistrarEncodedProto.

@Test
public void testRegistrarEncodedProto() throws Exception {
    PCollection<String> output = p.apply(testStream);
    AppliedPTransform<PBegin, PCollection<String>, TestStream<String>> appliedTestStream = AppliedPTransform.of("fakeName", PValues.expandInput(PBegin.in(p)), PValues.expandOutput(output), testStream, ResourceHints.create(), p);
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    RunnerApi.FunctionSpec spec = PTransformTranslation.toProto(appliedTestStream, components).getSpec();
    assertThat(spec.getUrn(), equalTo(TEST_STREAM_TRANSFORM_URN));
    RunnerApi.TestStreamPayload payload = TestStreamPayload.parseFrom(spec.getPayload());
    verifyTestStreamEncoding(testStream, payload, RehydratedComponents.forComponents(components.toComponents()));
}

Example 82 with Components

use of org.apache.beam.model.pipeline.v1.RunnerApi.Components in project beam by apache.

the class PTransformTranslationTest method toAndFromProto.

@Test
public void toAndFromProto() throws IOException {
    SdkComponents components = SdkComponents.create(spec.getTransform().getPipeline().getOptions());
    RunnerApi.PTransform converted = convert(spec, components);
    Components protoComponents = components.toComponents();
    // Sanity checks
    assertThat(converted.getInputsCount(), equalTo(spec.getTransform().getInputs().size()));
    assertThat(converted.getOutputsCount(), equalTo(spec.getTransform().getOutputs().size()));
    assertThat(converted.getSubtransformsCount(), equalTo(spec.getChildren().size()));
    assertThat(converted.getUniqueName(), equalTo(spec.getTransform().getFullName()));
    for (PValue inputValue : spec.getTransform().getInputs().values()) {
        PCollection<?> inputPc = (PCollection<?>) inputValue;
        protoComponents.getPcollectionsOrThrow(components.registerPCollection(inputPc));
    }
    for (PValue outputValue : spec.getTransform().getOutputs().values()) {
        PCollection<?> outputPc = (PCollection<?>) outputValue;
        protoComponents.getPcollectionsOrThrow(components.registerPCollection(outputPc));
    }
}

Example 83 with Components

use of org.apache.beam.model.pipeline.v1.RunnerApi.Components 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 84 with Components

use of org.apache.beam.model.pipeline.v1.RunnerApi.Components 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)));
}

Example 85 with Components

use of org.apache.beam.model.pipeline.v1.RunnerApi.Components in project beam by apache.

the class CreatePCollectionViewTranslationTest method testExtractionDirectFromTransform.

@Test
public void testExtractionDirectFromTransform() throws Exception {
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    components.registerPCollection(testPCollection);
    AppliedPTransform<?, ?, ?> appliedPTransform = AppliedPTransform.of("foo", PValues.expandInput(testPCollection), PValues.expandOutput(createViewTransform.getView()), createViewTransform, ResourceHints.create(), p);
    CreatePCollectionViewTranslation.getView((AppliedPTransform) appliedPTransform);
    FunctionSpec payload = PTransformTranslation.toProto(appliedPTransform, components).getSpec();
    // Checks that the payload is what it should be
    PCollectionView<?> deserializedView = (PCollectionView<?>) SerializableUtils.deserializeFromByteArray(payload.getPayload().toByteArray(), PCollectionView.class.getSimpleName());
    assertThat(deserializedView, Matchers.equalTo(createViewTransform.getView()));
}