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

Example 6 with PTransformMatcher

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

the class PTransformMatchers method splittableParDoSingle.

/**
   * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
   * that is splittable, as signified by {@link ProcessElementMethod#isSplittable()}.
   */
public static PTransformMatcher splittableParDoSingle() {
    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.processElement().isSplittable();
            }
            return false;
        }

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

Example 7 with PTransformMatcher

use of org.apache.beam.sdk.runners.PTransformMatcher 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 8 with PTransformMatcher

use of org.apache.beam.sdk.runners.PTransformMatcher 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 9 with PTransformMatcher

use of org.apache.beam.sdk.runners.PTransformMatcher 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 10 with PTransformMatcher

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

the class PTransformMatchersTest method parDoWithFnTypeWithMatchingType.

@Test
public void parDoWithFnTypeWithMatchingType() {
    DoFn<Object, Object> fn = new DoFn<Object, Object>() {

        @ProcessElement
        public void process(ProcessContext ctxt) {
        }
    };
    AppliedPTransform<?, ?, ?> parDoSingle = getAppliedTransform(ParDo.of(fn));
    AppliedPTransform<?, ?, ?> parDoMulti = getAppliedTransform(ParDo.of(fn).withOutputTags(new TupleTag<Object>(), TupleTagList.empty()));
    PTransformMatcher matcher = PTransformMatchers.parDoWithFnType(fn.getClass());
    assertThat(matcher.matches(parDoSingle), is(true));
    assertThat(matcher.matches(parDoMulti), is(true));
}