Source:
static void calculateJoinInputPartitions(JobGraph jobGraph, Config config) { Multimap<OperatorSpec, StreamEdge> joinSpecToStreamEdges = HashMultimap.create(); Multimap<StreamEdge, OperatorSpec> streamEdgeToJoinSpecs = HashMultimap.create(); Queue<OperatorSpec> joinQ = new LinkedList<>(); Set<OperatorSpec> visited = new HashSet<>(); StreamConfig streamConfig = new StreamConfig(config); jobGraph.getSpecGraph().getInputOperators().forEach((key, value) -> { StreamEdge streamEdge = jobGraph.getOrCreateStreamEdge(getStreamSpec(key, streamConfig)); findReachableJoins(value, streamEdge, joinSpecToStreamEdges, streamEdgeToJoinSpecs, joinQ, visited); }); while (!joinQ.isEmpty()) { OperatorSpec join = joinQ.poll(); int partitions = StreamEdge.PARTITIONS_UNKNOWN; for (StreamEdge edge : joinSpecToStreamEdges.get(join)) { int edgePartitions = edge.getPartitionCount(); if (edgePartitions != StreamEdge.PARTITIONS_UNKNOWN) { if (partitions == StreamEdge.PARTITIONS_UNKNOWN) { partitions = edgePartitions; } else if (partitions != edgePartitions) { throw new SamzaException(String.format( "Unable to resolve input partitions of stream %s for join. Expected: %d, Actual: %d", edge.getName(), partitions, edgePartitions)); } } } log.info("Inferred the partition count for the operator as {}.", join.getOpId(), partitions); for (StreamEdge edge : joinSpecToStreamEdges.get(join)) { if (edge.getPartitionCount() <= 0) { log.info("Initialize the partition count for the associated edge to {}.", edge.getName(), partitions); edge.setPartitionCount(partitions); for (OperatorSpec op : streamEdgeToJoinSpecs.get(edge)) { if (!visited.contains(op)) { joinQ.add(op); visited.add(op); } } } } } }

Target:
static void calculateJoinInputPartitions(JobGraph jobGraph, Config config) { Multimap<OperatorSpec, StreamEdge> joinSpecToStreamEdges = HashMultimap.create(); Multimap<StreamEdge, OperatorSpec> streamEdgeToJoinSpecs = HashMultimap.create(); Queue<OperatorSpec> joinQ = new LinkedList<>(); Set<OperatorSpec> visited = new HashSet<>(); StreamConfig streamConfig = new StreamConfig(config); jobGraph.getSpecGraph().getInputOperators().forEach((key, value) -> { StreamEdge streamEdge = jobGraph.getOrCreateStreamEdge(getStreamSpec(key, streamConfig)); findReachableJoins(value, streamEdge, joinSpecToStreamEdges, streamEdgeToJoinSpecs, joinQ, visited); }); while (!joinQ.isEmpty()) { OperatorSpec join = joinQ.poll(); int partitions = StreamEdge.PARTITIONS_UNKNOWN; for (StreamEdge edge : joinSpecToStreamEdges.get(join)) { int edgePartitions = edge.getPartitionCount(); if (edgePartitions != StreamEdge.PARTITIONS_UNKNOWN) { if (partitions == StreamEdge.PARTITIONS_UNKNOWN) { partitions = edgePartitions; log.info("Inferred the partition count for the join operator from {}." , new Object[] {partitions, join.getOpId(), edge.getName()}); } else if (partitions != edgePartitions) { throw new SamzaException(String.format( "Unable to resolve input partitions of stream %s for the join %s. Expected: %d, Actual: %d", edge.getName(), join.getOpId(), partitions, edgePartitions)); } } } for (StreamEdge edge : joinSpecToStreamEdges.get(join)) { if (edge.getPartitionCount() <= 0) { log.info("Set the partition count to for input stream to the join {}.", new Object[] {partitions, edge.getName(), join.getOpId()}); edge.setPartitionCount(partitions); for (OperatorSpec op : streamEdgeToJoinSpecs.get(edge)) { if (!visited.contains(op)) { joinQ.add(op); visited.add(op); } } } } } }