beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From tg...@apache.org
Subject [2/4] beam git commit: Add ParDoWithType PTransform Matcher
Date Fri, 17 Feb 2017 22:21:23 GMT
Add ParDoWithType PTransform Matcher

This matcher matches a ParDo where the contained DoFn matches a provided
type.


Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/e75fbe70
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/e75fbe70
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/e75fbe70

Branch: refs/heads/master
Commit: e75fbe70989ecfac41bb05fecc1dcd602d2a0b5c
Parents: de7b1a1
Author: Thomas Groh <tgroh@google.com>
Authored: Thu Feb 16 09:12:24 2017 -0800
Committer: Thomas Groh <tgroh@google.com>
Committed: Fri Feb 17 14:21:10 2017 -0800

----------------------------------------------------------------------
 .../core/construction/PTransformMatchers.java   | 21 ++++++++++
 .../construction/PTransformMatchersTest.java    | 43 ++++++++++++++++++++
 2 files changed, 64 insertions(+)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/e75fbe70/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
index 579ab7c..7b05ed1 100644
--- a/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
+++ b/runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java
@@ -162,6 +162,27 @@ public class PTransformMatchers {
   }
 
   /**
+   * A {@link PTransformMatcher} which matches a {@link ParDo.Bound} or {@link ParDo.BoundMulti}
+   * where the {@link DoFn} is of the provided type.
+   */
+  public static PTransformMatcher parDoWithFnType(final Class<? extends DoFn> fnType)
{
+    return new PTransformMatcher() {
+      @Override
+      public boolean matches(AppliedPTransform<?, ?, ?> application) {
+        DoFn<?, ?> fn;
+        if (application.getTransform() instanceof ParDo.Bound) {
+          fn = ((ParDo.Bound) application.getTransform()).getFn();
+        } else if (application.getTransform() instanceof ParDo.BoundMulti) {
+          fn = ((ParDo.BoundMulti) application.getTransform()).getFn();
+        } else {
+          return false;
+        }
+        return fnType.equals(fn.getClass());
+      }
+    };
+  }
+
+  /**
    * A {@link PTransformMatcher} which matches a {@link Flatten.FlattenPCollectionList} which
    * consumes no input {@link PCollection PCollections}.
    */

http://git-wip-us.apache.org/repos/asf/beam/blob/e75fbe70/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
----------------------------------------------------------------------
diff --git a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
index 1260289..439a475 100644
--- a/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
+++ b/runners/core-construction-java/src/test/java/org/apache/beam/runners/core/construction/PTransformMatchersTest.java
@@ -27,9 +27,11 @@ import com.google.common.base.MoreObjects;
 import java.io.Serializable;
 import java.util.Collections;
 import org.apache.beam.sdk.coders.VarIntCoder;
+import org.apache.beam.sdk.coders.VoidCoder;
 import org.apache.beam.sdk.runners.PTransformMatcher;
 import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.AppliedPTransform;
+import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.PTransform;
@@ -276,6 +278,47 @@ public class PTransformMatchersTest implements Serializable {
   }
 
   @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));
+  }
+
+  @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));
+  }
+
+  @Test
+  public void parDoWithFnTypeNotParDo() {
+    AppliedPTransform<?, ?, ?> notParDo = getAppliedTransform(Create.empty(VoidCoder.of()));
+    PTransformMatcher matcher = PTransformMatchers.parDoWithFnType(doFnWithState.getClass());
+    assertThat(matcher.matches(notParDo), is(false));
+  }
+
+  @Test
   public void emptyFlattenWithEmptyFlatten() {
     AppliedPTransform application =
         AppliedPTransform


Mime
View raw message