SAMZA-1221, SAMZA-1101: Internal cleanup for High-Level API implementation.

SAMZA-1221: Separated the OperatorSpec and MessageStream DAGs so that they're now duals of each other. Users interact with and construct the MessageStream DAG; we create and use the OperatorSpec DAG internally.
Moved common OperatorSpec functionality (getId, getOpCode, getOpName etc.) to the OperatorSpec abstract base class.
Added a new JoinOperatorSpec and PartialJoinOperatorImpls which are created from JoinOperatorSpec in OperatorGraphImpl.
Added a new InputOperatorSpec and InputOperatorImpl (previously RootOperatorImpl). InputOperatorSpec is created when StreamGraph#getInputStream is called.
SAMZA-1101: Added a new OutputOperatorSpec and OutputOperatorImpl for partitionBy and sendTo. These are Separate from SinkOperatorSpec for and SinkOperatorImpl for sink. We don't need to create a sinkFn for partitionBy and sendTo anymore.
Updated most unit tests to use the new classes and avoid reflection.

Author: Prateek Maheshwari <pmaheshw@linkedin.com>

Reviewers: Jagadish V <vjagadish1989@apache.org>

Closes apache#194 from prateekm/internal-cleanup

Author: prateekm

.gitignore

@@ -28,3 +28,4 @@ docs/learn/documentation/*/rest/javadocs
 out/
 *.patch
 **.pyc
+samza-shell/src/main/visualizer/plan.json

samza-core/src/main/java/org/apache/samza/execution/ExecutionPlanner.java

@@ -31,11 +31,9 @@
 import org.apache.samza.SamzaException;
 import org.apache.samza.config.Config;
 import org.apache.samza.config.JobConfig;
-import org.apache.samza.operators.MessageStream;
-import org.apache.samza.operators.MessageStreamImpl;
 import org.apache.samza.operators.StreamGraphImpl;
+import org.apache.samza.operators.spec.JoinOperatorSpec;
 import org.apache.samza.operators.spec.OperatorSpec;
-import org.apache.samza.operators.spec.PartialJoinOperatorSpec;
 import org.apache.samza.system.StreamSpec;
 import org.apache.samza.system.SystemStream;
 import org.slf4j.Logger;
@@ -77,7 +75,7 @@ public ExecutionPlan plan(StreamGraphImpl streamGraph) throws Exception {
    */
   /* package private */ JobGraph createJobGraph(StreamGraphImpl streamGraph) {
     JobGraph jobGraph = new JobGraph(config);
-    Set<StreamSpec> sourceStreams = new HashSet<>(streamGraph.getInputStreams().keySet());
+    Set<StreamSpec> sourceStreams = new HashSet<>(streamGraph.getInputOperators().keySet());
     Set<StreamSpec> sinkStreams = new HashSet<>(streamGraph.getOutputStreams().keySet());
     Set<StreamSpec> intStreams = new HashSet<>(sourceStreams);
     intStreams.retainAll(sinkStreams);
@@ -120,7 +118,7 @@ public ExecutionPlan plan(StreamGraphImpl streamGraph) throws Exception {
   /**
    * Fetch the partitions of source/sink streams and update the StreamEdges.
    * @param jobGraph {@link JobGraph}
-   * @param streamManager the {@StreamManager} to interface with the streams.
+   * @param streamManager the {@link StreamManager} to interface with the streams.
    */
   /* package private */ static void updateExistingPartitions(JobGraph jobGraph, StreamManager streamManager) {
     Set<StreamEdge> existingStreams = new HashSet<>();
@@ -157,20 +155,16 @@ public ExecutionPlan plan(StreamGraphImpl streamGraph) throws Exception {
     Multimap<OperatorSpec, StreamEdge> joinSpecToStreamEdges = HashMultimap.create();
     // reverse mapping of the above
     Multimap<StreamEdge, OperatorSpec> streamEdgeToJoinSpecs = HashMultimap.create();
-    // Mapping from the output stream to the join spec. Since StreamGraph creates two partial join operators for a join and they
-    // will have the same output stream, this mapping is used to choose one of them as the unique join spec representing this join
-    // (who register first in the map wins).
-    Map<MessageStream, OperatorSpec> outputStreamToJoinSpec = new HashMap<>();
     // A queue of joins with known input partitions
     Queue<OperatorSpec> joinQ = new LinkedList<>();
     // The visited set keeps track of the join specs that have been already inserted in the queue before
     Set<OperatorSpec> visited = new HashSet<>();
 
-    streamGraph.getInputStreams().entrySet().forEach(entry -> {
+    streamGraph.getInputOperators().entrySet().forEach(entry -> {
         StreamEdge streamEdge = jobGraph.getOrCreateStreamEdge(entry.getKey());
         // Traverses the StreamGraph to find and update mappings for all Joins reachable from this input StreamEdge
         findReachableJoins(entry.getValue(), streamEdge, joinSpecToStreamEdges, streamEdgeToJoinSpecs,
-            outputStreamToJoinSpec, joinQ, visited);
+            joinQ, visited);
       });
 
     // At this point, joinQ contains joinSpecs where at least one of the input stream edge partitions is known.
@@ -209,44 +203,33 @@ public ExecutionPlan plan(StreamGraphImpl streamGraph) throws Exception {
   }
 
   /**
-   * This function traverses the StreamGraph to find and update mappings for all Joins reachable from this input StreamEdge
-   * @param inputMessageStream next input MessageStream to traverse {@link MessageStream}
+   * This function traverses the {@link OperatorSpec} graph to find and update mappings for all Joins reachable
+   * from this input {@link StreamEdge}.
+   * @param operatorSpec the {@link OperatorSpec} to traverse
    * @param sourceStreamEdge source {@link StreamEdge}
    * @param joinSpecToStreamEdges mapping from join spec to its source {@link StreamEdge}s
    * @param streamEdgeToJoinSpecs mapping from source {@link StreamEdge} to the join specs that consumes it
-   * @param outputStreamToJoinSpec mapping from the output stream to the join spec
    * @param joinQ queue that contains joinSpecs where at least one of the input stream edge partitions is known.
    */
-  private static void findReachableJoins(MessageStream inputMessageStream, StreamEdge sourceStreamEdge,
-      Multimap<OperatorSpec, StreamEdge> joinSpecToStreamEdges, Multimap<StreamEdge, OperatorSpec> streamEdgeToJoinSpecs,
-      Map<MessageStream, OperatorSpec> outputStreamToJoinSpec, Queue<OperatorSpec> joinQ, Set<OperatorSpec> visited) {
-    Collection<OperatorSpec> specs = ((MessageStreamImpl) inputMessageStream).getRegisteredOperatorSpecs();
-    for (OperatorSpec spec : specs) {
-      if (spec instanceof PartialJoinOperatorSpec) {
-        // every join will have two partial join operators
-        // we will choose one of them in order to consolidate the inputs
-        // the first one who registered with the outputStreamToJoinSpec will win
-        MessageStream output = spec.getNextStream();
-        OperatorSpec joinSpec = outputStreamToJoinSpec.get(output);
-        if (joinSpec == null) {
-          joinSpec = spec;
-          outputStreamToJoinSpec.put(output, joinSpec);
-        }
-
-        joinSpecToStreamEdges.put(joinSpec, sourceStreamEdge);
-        streamEdgeToJoinSpecs.put(sourceStreamEdge, joinSpec);
-
-        if (!visited.contains(joinSpec) && sourceStreamEdge.getPartitionCount() > 0) {
-          // put the joins with known input partitions into the queue
-          joinQ.add(joinSpec);
-          visited.add(joinSpec);
-        }
+  private static void findReachableJoins(OperatorSpec operatorSpec, StreamEdge sourceStreamEdge,
+      Multimap<OperatorSpec, StreamEdge> joinSpecToStreamEdges,
+      Multimap<StreamEdge, OperatorSpec> streamEdgeToJoinSpecs,
+      Queue<OperatorSpec> joinQ, Set<OperatorSpec> visited) {
+    if (operatorSpec instanceof JoinOperatorSpec) {
+      joinSpecToStreamEdges.put(operatorSpec, sourceStreamEdge);
+      streamEdgeToJoinSpecs.put(sourceStreamEdge, operatorSpec);
+
+      if (!visited.contains(operatorSpec) && sourceStreamEdge.getPartitionCount() > 0) {
+        // put the joins with known input partitions into the queue and mark as visited
+        joinQ.add(operatorSpec);
+        visited.add(operatorSpec);
       }
+    }
 
-      if (spec.getNextStream() != null) {
-        findReachableJoins(spec.getNextStream(), sourceStreamEdge, joinSpecToStreamEdges, streamEdgeToJoinSpecs, outputStreamToJoinSpec, joinQ,
-            visited);
-      }
+    Collection<OperatorSpec> registeredOperatorSpecs = operatorSpec.getRegisteredOperatorSpecs();
+    for (OperatorSpec registeredOpSpec : registeredOperatorSpecs) {
+      findReachableJoins(registeredOpSpec, sourceStreamEdge, joinSpecToStreamEdges, streamEdgeToJoinSpecs, joinQ,
+          visited);
     }
   }
 

samza-core/src/main/java/org/apache/samza/execution/JobGraphJsonGenerator.java

@@ -19,24 +19,21 @@
 
 package org.apache.samza.execution;
 
-import com.google.common.base.Joiner;
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.Multimap;
 import java.io.ByteArrayOutputStream;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.apache.samza.config.ApplicationConfig;
-import org.apache.samza.operators.MessageStream;
-import org.apache.samza.operators.MessageStreamImpl;
+import org.apache.samza.operators.spec.JoinOperatorSpec;
 import org.apache.samza.operators.spec.OperatorSpec;
-import org.apache.samza.operators.util.OperatorJsonUtils;
+import org.apache.samza.operators.spec.OperatorSpec.OpCode;
+import org.apache.samza.operators.spec.OutputOperatorSpec;
+import org.apache.samza.operators.spec.OutputStreamImpl;
 import org.codehaus.jackson.annotate.JsonProperty;
 import org.codehaus.jackson.map.ObjectMapper;
 
@@ -73,8 +70,6 @@ static final class OperatorGraphJson {
     List<StreamJson> outputStreams;
     @JsonProperty("operators")
     Map<Integer, Map<String, Object>> operators = new HashMap<>();
-    @JsonProperty("canonicalOpIds")
-    Map<Integer, String> canonicalOpIds = new HashMap<>();
   }
 
   static final class StreamJson {
@@ -108,11 +103,6 @@ static final class JobGraphJson {
     String applicationId;
   }
 
-  // Mapping from the output stream to the ids.
-  // Logically they belong to the same operator, but in code we generate one operator for each input.
-  // This is to associate the operators that output to the same MessageStream.
-  Multimap<MessageStream, Integer> outputStreamToOpIds = HashMultimap.create();
-
   /**
    * Returns the JSON representation of a {@link JobGraph}
    * @param jobGraph {@link JobGraph}
@@ -157,28 +147,21 @@ private JobNodeJson buildJobNodeJson(JobNode jobNode) {
   }
 
   /**
-   * Traverse the {@StreamGraph} and build the operator graph JSON POJO.
+   * Traverse the {@link OperatorSpec} graph and build the operator graph JSON POJO.
    * @param jobNode job node in the {@link JobGraph}
    * @return {@link org.apache.samza.execution.JobGraphJsonGenerator.OperatorGraphJson}
    */
   private OperatorGraphJson buildOperatorGraphJson(JobNode jobNode) {
     OperatorGraphJson opGraph = new OperatorGraphJson();
     opGraph.inputStreams = new ArrayList<>();
-    jobNode.getStreamGraph().getInputStreams().forEach((streamSpec, stream) -> {
+    jobNode.getStreamGraph().getInputOperators().forEach((streamSpec, operatorSpec) -> {
         StreamJson inputJson = new StreamJson();
         opGraph.inputStreams.add(inputJson);
         inputJson.streamId = streamSpec.getId();
-        Collection<OperatorSpec> specs = ((MessageStreamImpl) stream).getRegisteredOperatorSpecs();
+        Collection<OperatorSpec> specs = operatorSpec.getRegisteredOperatorSpecs();
         inputJson.nextOperatorIds = specs.stream().map(OperatorSpec::getOpId).collect(Collectors.toSet());
 
-        updateOperatorGraphJson((MessageStreamImpl) stream, opGraph);
-
-        for (Map.Entry<MessageStream, Collection<Integer>> entry : outputStreamToOpIds.asMap().entrySet()) {
-          List<Integer> sortedIds = new ArrayList<>(entry.getValue());
-          Collections.sort(sortedIds);
-          String canonicalId = Joiner.on(',').join(sortedIds);
-          sortedIds.stream().forEach(id -> opGraph.canonicalOpIds.put(id, canonicalId));
-        }
+        updateOperatorGraphJson(operatorSpec, opGraph);
       });
 
     opGraph.outputStreams = new ArrayList<>();
@@ -191,23 +174,43 @@ private OperatorGraphJson buildOperatorGraphJson(JobNode jobNode) {
   }
 
   /**
-   * Traverse the {@StreamGraph} recursively and update the operator graph JSON POJO.
-   * @param messageStream input
+   * Traverse the {@link OperatorSpec} graph recursively and update the operator graph JSON POJO.
+   * @param operatorSpec input
    * @param opGraph operator graph to build
    */
-  private void updateOperatorGraphJson(MessageStreamImpl messageStream, OperatorGraphJson opGraph) {
-    Collection<OperatorSpec> specs = messageStream.getRegisteredOperatorSpecs();
-    specs.forEach(opSpec -> {
-        opGraph.operators.put(opSpec.getOpId(), OperatorJsonUtils.operatorToMap(opSpec));
-
-        if (opSpec.getOpCode() == OperatorSpec.OpCode.JOIN || opSpec.getOpCode() == OperatorSpec.OpCode.MERGE) {
-          outputStreamToOpIds.put(opSpec.getNextStream(), opSpec.getOpId());
-        }
-
-        if (opSpec.getNextStream() != null) {
-          updateOperatorGraphJson(opSpec.getNextStream(), opGraph);
-        }
-      });
+  private void updateOperatorGraphJson(OperatorSpec operatorSpec, OperatorGraphJson opGraph) {
+    // TODO xiliu: render input operators instead of input streams
+    if (operatorSpec.getOpCode() != OpCode.INPUT) {
+      opGraph.operators.put(operatorSpec.getOpId(), operatorToMap(operatorSpec));
+    }
+    Collection<OperatorSpec> specs = operatorSpec.getRegisteredOperatorSpecs();
+    specs.forEach(opSpec -> updateOperatorGraphJson(opSpec, opGraph));
+  }
+
+  /**
+   * Format the operator properties into a map
+   * @param spec a {@link OperatorSpec} instance
+   * @return map of the operator properties
+   */
+  private Map<String, Object> operatorToMap(OperatorSpec spec) {
+    Map<String, Object> map = new HashMap<>();
+    map.put("opCode", spec.getOpCode().name());
+    map.put("opId", spec.getOpId());
+    map.put("sourceLocation", spec.getSourceLocation());
+
+    Collection<OperatorSpec> nextOperators = spec.getRegisteredOperatorSpecs();
+    map.put("nextOperatorIds", nextOperators.stream().map(OperatorSpec::getOpId).collect(Collectors.toSet()));
+
+    if (spec instanceof OutputOperatorSpec) {
+      OutputStreamImpl outputStream = ((OutputOperatorSpec) spec).getOutputStream();
+      map.put("outputStreamId", outputStream.getStreamSpec().getId());
+    }
+
+    if (spec instanceof JoinOperatorSpec) {
+      map.put("ttlMs", ((JoinOperatorSpec) spec).getTtlMs());
+    }
+
+    return map;
   }
 
   /**

samza-core/src/main/java/org/apache/samza/execution/JobNode.java

@@ -31,8 +31,8 @@
 import org.apache.samza.config.MapConfig;
 import org.apache.samza.config.TaskConfig;
 import org.apache.samza.operators.StreamGraphImpl;
+import org.apache.samza.operators.spec.JoinOperatorSpec;
 import org.apache.samza.operators.spec.OperatorSpec;
-import org.apache.samza.operators.spec.PartialJoinOperatorSpec;
 import org.apache.samza.operators.spec.WindowOperatorSpec;
 import org.apache.samza.operators.util.MathUtils;
 import org.apache.samza.util.Util;
@@ -145,7 +145,7 @@ private long computeTriggerInterval() {
     // Filter out the join operators, and obtain a list of their ttl values
     List<Long> joinTtlIntervals = operatorSpecs.stream()
         .filter(spec -> spec.getOpCode() == OperatorSpec.OpCode.JOIN)
-        .map(spec -> ((PartialJoinOperatorSpec) spec).getTtlMs())
+        .map(spec -> ((JoinOperatorSpec) spec).getTtlMs())
         .collect(Collectors.toList());
 
     // Combine both the above lists