Examples with AppliedPTransform org.apache.beam.sdk.runners.AppliedPTransform used on opensource projects

Example 11 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class PTransformMatchers method splittableParDoMulti.

/**
   * A {@link PTransformMatcher} that matches a {@link ParDo.MultiOutput} containing a {@link DoFn}
   * that is splittable, as signified by {@link ProcessElementMethod#isSplittable()}.
   */
public static PTransformMatcher splittableParDoMulti() {
    return new PTransformMatcher() {

        @Override
        public boolean matches(AppliedPTransform<?, ?, ?> application) {
            PTransform<?, ?> transform = application.getTransform();
            if (transform instanceof ParDo.MultiOutput) {
                DoFn<?, ?> fn = ((ParDo.MultiOutput<?, ?>) transform).getFn();
                DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
                return signature.processElement().isSplittable();
            }
            return false;
        }

        @Override
        public String toString() {
            return MoreObjects.toStringHelper("SplittableParDoMultiMatcher").toString();
        }
    };
}

Example 12 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class PTransformMatchers method stateOrTimerParDoSingle.

/**
   * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
   * that uses state or timers, as specified by {@link DoFnSignature#usesState()} and {@link
   * DoFnSignature#usesTimers()}.
   */
public static PTransformMatcher stateOrTimerParDoSingle() {
    return new PTransformMatcher() {

        @Override
        public boolean matches(AppliedPTransform<?, ?, ?> application) {
            PTransform<?, ?> transform = application.getTransform();
            if (transform instanceof ParDo.SingleOutput) {
                DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn();
                DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
                return signature.usesState() || signature.usesTimers();
            }
            return false;
        }

        @Override
        public String toString() {
            return MoreObjects.toStringHelper("StateOrTimerParDoSingleMatcher").toString();
        }
    };
}

Example 13 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class WatermarkManager method refreshWatermarks.

private Set<AppliedPTransform<?, ?, ?>> refreshWatermarks(AppliedPTransform<?, ?, ?> toRefresh) {
    TransformWatermarks myWatermarks = transformToWatermarks.get(toRefresh);
    WatermarkUpdate updateResult = myWatermarks.refresh();
    if (updateResult.isAdvanced()) {
        Set<AppliedPTransform<?, ?, ?>> additionalRefreshes = new HashSet<>();
        for (PValue outputPValue : toRefresh.getOutputs().values()) {
            additionalRefreshes.addAll(graph.getPrimitiveConsumers(outputPValue));
        }
        return additionalRefreshes;
    }
    return Collections.emptySet();
}

Example 14 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform in project beam by apache.

the class PTransformMatchers method stateOrTimerParDoMulti.

/**
   * A {@link PTransformMatcher} that matches a {@link ParDo.MultiOutput} containing a {@link DoFn}
   * that uses state or timers, as specified by {@link DoFnSignature#usesState()} and
   * {@link DoFnSignature#usesTimers()}.
   */
public static PTransformMatcher stateOrTimerParDoMulti() {
    return new PTransformMatcher() {

        @Override
        public boolean matches(AppliedPTransform<?, ?, ?> application) {
            PTransform<?, ?> transform = application.getTransform();
            if (transform instanceof ParDo.MultiOutput) {
                DoFn<?, ?> fn = ((ParDo.MultiOutput<?, ?>) transform).getFn();
                DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
                return signature.usesState() || signature.usesTimers();
            }
            return false;
        }

        @Override
        public String toString() {
            return MoreObjects.toStringHelper("StateOrTimerParDoMultiMatcher").toString();
        }
    };
}

Example 15 with AppliedPTransform

use of org.apache.beam.sdk.runners.AppliedPTransform 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();
}