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

Example 1 with PTransformMatcher

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

the class PTransformMatchersTest method classEqualToMatchesSameClass.

@Test
public void classEqualToMatchesSameClass() {
    PTransformMatcher matcher = PTransformMatchers.classEqualTo(ParDo.SingleOutput.class);
    AppliedPTransform<?, ?, ?> application = getAppliedTransform(ParDo.of(new DoFn<KV<String, Integer>, Integer>() {

        @ProcessElement
        public void doStuff(ProcessContext ctxt) {
        }
    }));
    assertThat(matcher.matches(application), is(true));
}

Example 2 with PTransformMatcher

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

the class PTransformMatchersTest method parDoWithFnTypeWithNoMatch.

@Test
public void parDoWithFnTypeWithNoMatch() {
    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(doFnWithState.getClass());
    assertThat(matcher.matches(parDoSingle), is(false));
    assertThat(matcher.matches(parDoMulti), is(false));
}

Example 3 with PTransformMatcher

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

the class PTransformMatchersTest method classEqualToDoesNotMatchSubclass.

@Test
public void classEqualToDoesNotMatchSubclass() {
    class MyPTransform extends PTransform<PCollection<KV<String, Integer>>, PCollection<Integer>> {

        @Override
        public PCollection<Integer> expand(PCollection<KV<String, Integer>> input) {
            return PCollection.createPrimitiveOutputInternal(input.getPipeline(), input.getWindowingStrategy(), input.isBounded());
        }
    }
    PTransformMatcher matcher = PTransformMatchers.classEqualTo(MyPTransform.class);
    MyPTransform subclass = new MyPTransform() {
    };
    assertThat(subclass.getClass(), not(Matchers.<Class<?>>equalTo(MyPTransform.class)));
    assertThat(subclass, instanceOf(MyPTransform.class));
    AppliedPTransform<?, ?, ?> application = getAppliedTransform(subclass);
    assertThat(matcher.matches(application), is(false));
}

Example 4 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 5 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();
        }
    };
}