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

Example 6 with FunctionSpec

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

the class ExecutableStage method fromPayload.

/**
 * Return an {@link ExecutableStage} constructed from the provided {@link FunctionSpec}
 * representation.
 *
 * <p>See {@link #toPTransform} for how the payload is constructed.
 *
 * <p>Note: The payload contains some information redundant with the {@link PTransform} it is the
 * payload of. The {@link ExecutableStagePayload} should be sufficiently rich to construct a
 * {@code ProcessBundleDescriptor} using only the payload.
 */
static ExecutableStage fromPayload(ExecutableStagePayload payload) {
    Components components = payload.getComponents();
    Environment environment = payload.getEnvironment();
    Collection<WireCoderSetting> wireCoderSettings = payload.getWireCoderSettingsList();
    PCollectionNode input = PipelineNode.pCollection(payload.getInput(), components.getPcollectionsOrThrow(payload.getInput()));
    List<SideInputReference> sideInputs = payload.getSideInputsList().stream().map(sideInputId -> SideInputReference.fromSideInputId(sideInputId, components)).collect(Collectors.toList());
    List<UserStateReference> userStates = payload.getUserStatesList().stream().map(userStateId -> UserStateReference.fromUserStateId(userStateId, components)).collect(Collectors.toList());
    List<TimerReference> timers = payload.getTimersList().stream().map(timerId -> TimerReference.fromTimerId(timerId, components)).collect(Collectors.toList());
    List<PTransformNode> transforms = payload.getTransformsList().stream().map(id -> PipelineNode.pTransform(id, components.getTransformsOrThrow(id))).collect(Collectors.toList());
    List<PCollectionNode> outputs = payload.getOutputsList().stream().map(id -> PipelineNode.pCollection(id, components.getPcollectionsOrThrow(id))).collect(Collectors.toList());
    return ImmutableExecutableStage.of(components, environment, input, sideInputs, userStates, timers, transforms, outputs, wireCoderSettings);
}

Example 7 with FunctionSpec

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

the class CreatePCollectionViewTranslationTest method testEncodedProto.

@Test
public void testEncodedProto() 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);
    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()));
}

Example 8 with FunctionSpec

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

the class PTransformTranslation method toProto.

/**
   * Translates an {@link AppliedPTransform} into a runner API proto.
   *
   * <p>Does not register the {@code appliedPTransform} within the provided {@link SdkComponents}.
   */
static RunnerApi.PTransform toProto(AppliedPTransform<?, ?, ?> appliedPTransform, List<AppliedPTransform<?, ?, ?>> subtransforms, SdkComponents components) throws IOException {
    RunnerApi.PTransform.Builder transformBuilder = RunnerApi.PTransform.newBuilder();
    for (Map.Entry<TupleTag<?>, PValue> taggedInput : appliedPTransform.getInputs().entrySet()) {
        checkArgument(taggedInput.getValue() instanceof PCollection, "Unexpected input type %s", taggedInput.getValue().getClass());
        transformBuilder.putInputs(toProto(taggedInput.getKey()), components.registerPCollection((PCollection<?>) taggedInput.getValue()));
    }
    for (Map.Entry<TupleTag<?>, PValue> taggedOutput : appliedPTransform.getOutputs().entrySet()) {
        // TODO: Remove gating
        if (taggedOutput.getValue() instanceof PCollection) {
            checkArgument(taggedOutput.getValue() instanceof PCollection, "Unexpected output type %s", taggedOutput.getValue().getClass());
            transformBuilder.putOutputs(toProto(taggedOutput.getKey()), components.registerPCollection((PCollection<?>) taggedOutput.getValue()));
        }
    }
    for (AppliedPTransform<?, ?, ?> subtransform : subtransforms) {
        transformBuilder.addSubtransforms(components.getExistingPTransformId(subtransform));
    }
    transformBuilder.setUniqueName(appliedPTransform.getFullName());
    // TODO: Display Data
    PTransform<?, ?> transform = appliedPTransform.getTransform();
    if (KNOWN_PAYLOAD_TRANSLATORS.containsKey(transform.getClass())) {
        FunctionSpec payload = KNOWN_PAYLOAD_TRANSLATORS.get(transform.getClass()).translate(appliedPTransform, components);
        transformBuilder.setSpec(payload);
    }
    return transformBuilder.build();
}

Example 9 with FunctionSpec

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

the class ParDoTranslation method viewFnFromProto.

private static ViewFn<?, ?> viewFnFromProto(SdkFunctionSpec viewFn) throws InvalidProtocolBufferException {
    FunctionSpec spec = viewFn.getSpec();
    checkArgument(spec.getUrn().equals(CUSTOM_JAVA_VIEW_FN_URN), "Can't deserialize unknown %s type %s", ViewFn.class.getSimpleName(), spec.getUrn());
    return (ViewFn<?, ?>) SerializableUtils.deserializeFromByteArray(spec.getParameter().unpack(BytesValue.class).getValue().toByteArray(), "Custom ViewFn");
}

Example 10 with FunctionSpec

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

the class ParDoTranslation method windowMappingFnFromProto.

private static WindowMappingFn<?> windowMappingFnFromProto(SdkFunctionSpec windowMappingFn) throws InvalidProtocolBufferException {
    FunctionSpec spec = windowMappingFn.getSpec();
    checkArgument(spec.getUrn().equals(CUSTOM_JAVA_WINDOW_MAPPING_FN_URN), "Can't deserialize unknown %s type %s", WindowMappingFn.class.getSimpleName(), spec.getUrn());
    return (WindowMappingFn<?>) SerializableUtils.deserializeFromByteArray(spec.getParameter().unpack(BytesValue.class).getValue().toByteArray(), "Custom WinodwMappingFn");
}