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

Example 31 with PTransform

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

the class OutputDeduplicator method deduplicatePCollections.

private static PTransformDeduplication deduplicatePCollections(PTransformNode transform, Collection<PCollectionNode> duplicates, Predicate<String> existingPCollectionIds) {
    Map<String, PCollectionNode> unzippedOutputs = createPartialPCollections(duplicates, existingPCollectionIds);
    PTransform pTransform = updateOutputs(transform.getTransform(), unzippedOutputs);
    return PTransformDeduplication.of(PipelineNode.pTransform(transform.getId(), pTransform), unzippedOutputs);
}

Example 32 with PTransform

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

the class SideInputReference method fromSideInputId.

/**
 * Create a side input reference from a SideInputId proto and components.
 */
public static SideInputReference fromSideInputId(SideInputId sideInputId, RunnerApi.Components components) {
    String transformId = sideInputId.getTransformId();
    String localName = sideInputId.getLocalName();
    String collectionId = components.getTransformsOrThrow(transformId).getInputsOrThrow(localName);
    PTransform transform = components.getTransformsOrThrow(transformId);
    PCollection collection = components.getPcollectionsOrThrow(collectionId);
    return SideInputReference.of(PipelineNode.pTransform(transformId, transform), localName, PipelineNode.pCollection(collectionId, collection));
}

Example 33 with PTransform

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

the class RegisterAndProcessBundleOperation method extractCrossBoundaryGrpcPCollectionNames.

private Set<String> extractCrossBoundaryGrpcPCollectionNames(final Set<Entry<String, PTransform>> ptransforms) {
    Set<String> result = new HashSet<>();
    // GRPC Read/Write expected to only have one Output/Input respectively.
    for (Map.Entry<String, RunnerApi.PTransform> pTransform : ptransforms) {
        if (pTransform.getValue().getSpec().getUrn().equals(RemoteGrpcPortRead.URN)) {
            String grpcReadTransformOutputName = Iterables.getOnlyElement(pTransform.getValue().getOutputsMap().keySet());
            String pcollectionName = pTransform.getValue().getOutputsMap().get(grpcReadTransformOutputName);
            result.add(pcollectionName);
        }
        if (pTransform.getValue().getSpec().getUrn().equals(RemoteGrpcPortWrite.URN)) {
            String grpcTransformInputName = Iterables.getOnlyElement(pTransform.getValue().getInputsMap().keySet());
            String pcollectionName = pTransform.getValue().getInputsMap().get(grpcTransformInputName);
            result.add(pcollectionName);
        }
    }
    return result;
}

Example 34 with PTransform

use of org.apache.beam.sdk.common.runner.v1.RunnerApi.PTransform 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 35 with PTransform

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

the class ParDoTranslation method getMainInput.

public static RunnerApi.PCollection getMainInput(RunnerApi.PTransform ptransform, Components components) throws IOException {
    checkArgument(ptransform.getSpec().getUrn().equals(PAR_DO_TRANSFORM_URN), "Unexpected payload type %s", ptransform.getSpec().getUrn());
    ParDoPayload payload = ptransform.getSpec().getParameter().unpack(ParDoPayload.class);
    String mainInputId = Iterables.getOnlyElement(Sets.difference(ptransform.getInputsMap().keySet(), payload.getSideInputsMap().keySet()));
    return components.getPcollectionsOrThrow(ptransform.getInputsOrThrow(mainInputId));
}