[incubator-nemo] branch master updated: [NEMO-67] Fix Communication Pattern for Beam CoGroupByKey transform (#23)

[wo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wonook pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-nemo.git


The following commit(s) were added to refs/heads/master by this push:
     new 94761f4  [NEMO-67] Fix Communication Pattern for Beam CoGroupByKey transform (#23)
94761f4 is described below

commit 94761f42b1787b2b098af729355120098021546b
Author: JangHo Seo <ja...@jangho.io>
AuthorDate: Mon May 28 17:20:24 2018 +0900

    [NEMO-67] Fix Communication Pattern for Beam CoGroupByKey transform (#23)
    
    JIRA: [NEMO-67: Support Beam CoGroupByKey transform](https://issues.apache.org/jira/projects/NEMO/issues/NEMO-67)
    
    **Major changes:**
    - Modify NemoPipelineVisitor to fix communication pattern for Beam CoGroupByKey transform
    
    **Minor changes to note:**
    - None
    
    **Tests for the changes:**
    - Currently, none.
    - We can add example join application later by [NEMO-71: Add LocationShareAssignmentPass and Example Application](https://issues.apache.org/jira/projects/NEMO/issues/NEMO-71).
    
    **Other comments:**
    - None
    
    resolves [NEMO-67](https://issues.apache.org/jira/projects/NEMO/issues/NEMO-67)
---
 .../frontend/beam/NemoPipelineVisitor.java         | 27 ++++++++++++++++++----
 .../frontend/beam/transform/DoTransform.java       |  7 ++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/NemoPipelineVisitor.java b/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/NemoPipelineVisitor.java
index a649d72..f61adc6 100644
--- a/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/NemoPipelineVisitor.java
+++ b/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/NemoPipelineVisitor.java
@@ -15,6 +15,7 @@
  */
 package edu.snu.nemo.compiler.frontend.beam;
 
+import edu.snu.nemo.common.ir.vertex.transform.Transform;
 import edu.snu.nemo.compiler.frontend.beam.coder.BeamCoder;
 import edu.snu.nemo.common.dag.DAGBuilder;
 import edu.snu.nemo.common.ir.edge.IREdge;
@@ -241,13 +242,29 @@ public final class NemoPipelineVisitor extends Pipeline.PipelineVisitor.Defaults
    */
   private static DataCommunicationPatternProperty.Value getEdgeCommunicationPattern(final IRVertex src,
                                                                                     final IRVertex dst) {
-    if (dst instanceof OperatorVertex && ((OperatorVertex) dst).getTransform() instanceof GroupByKeyTransform) {
+    final Class<?> constructUnionTableFn;
+    try {
+      constructUnionTableFn = Class.forName("org.apache.beam.sdk.transforms.join.CoGroupByKey$ConstructUnionTableFn");
+    } catch (final ClassNotFoundException e) {
+      throw new RuntimeException(e);
+    }
+
+    final Transform srcTransform = src instanceof OperatorVertex ? ((OperatorVertex) src).getTransform() : null;
+    final Transform dstTransform = dst instanceof OperatorVertex ? ((OperatorVertex) dst).getTransform() : null;
+    final DoFn srcDoFn = srcTransform instanceof DoTransform ? ((DoTransform) srcTransform).getDoFn() : null;
+
+    if (srcDoFn != null && srcDoFn.getClass().equals(constructUnionTableFn)) {
       return DataCommunicationPatternProperty.Value.Shuffle;
-    } else if (dst instanceof OperatorVertex && ((OperatorVertex) dst).getTransform() instanceof CreateViewTransform
-        || src instanceof OperatorVertex && ((OperatorVertex) src).getTransform() instanceof CreateViewTransform) {
-      return DataCommunicationPatternProperty.Value.BroadCast;
-    } else {
+    }
+    if (srcTransform instanceof FlattenTransform) {
       return DataCommunicationPatternProperty.Value.OneToOne;
     }
+    if (dstTransform instanceof GroupByKeyTransform) {
+      return DataCommunicationPatternProperty.Value.Shuffle;
+    }
+    if (srcTransform instanceof CreateViewTransform || dstTransform instanceof CreateViewTransform) {
+      return DataCommunicationPatternProperty.Value.BroadCast;
+    }
+    return DataCommunicationPatternProperty.Value.OneToOne;
   }
 }
diff --git a/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/transform/DoTransform.java b/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/transform/DoTransform.java
index e099253..883d401 100644
--- a/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/transform/DoTransform.java
+++ b/compiler/frontend/beam/src/main/java/edu/snu/nemo/compiler/frontend/beam/transform/DoTransform.java
@@ -328,4 +328,11 @@ public final class DoTransform<I, O> implements Transform<I, O> {
       throw new UnsupportedOperationException("timer() in ProcessContext under DoTransform");
     }
   }
+
+  /**
+   * @return {@link DoFn} for this transform.
+   */
+  public DoFn getDoFn() {
+    return doFn;
+  }
 }

-- 
To stop receiving notification emails like this one, please contact
wonook@apache.org.