private static void setJoinedIntermediateStreamPartitions(Collection<StreamSet> joinedStreamSets) {
    // Map every intermediate stream to all the stream-sets it appears in
    Multimap<StreamEdge, StreamSet> intermediateStreamToStreamSets = HashMultimap.create();
    for (StreamSet streamSet : joinedStreamSets) {
      for (StreamEdge streamEdge : streamSet.getStreamEdges()) {
        if (streamEdge.getPartitionCount() == StreamEdge.PARTITIONS_UNKNOWN) {
          intermediateStreamToStreamSets.put(streamEdge, streamSet);
        }
      }
    }

    Set<StreamSet> streamSets = new HashSet<>(joinedStreamSets);
    Set<StreamSet> processedStreamSets = new HashSet<>();

    while (!streamSets.isEmpty()) {
      // Retrieve and remove one stream set
      StreamSet streamSet = streamSets.iterator().next();
      streamSets.remove(streamSet);

      // Find any stream with set partitions in this set
      Optional<StreamEdge> streamWithSetPartitions =
          streamSet.getStreamEdges().stream()
              .filter(streamEdge -> streamEdge.getPartitionCount() != StreamEdge.PARTITIONS_UNKNOWN)
              .findAny();

      if (streamWithSetPartitions.isPresent()) {
        // Mark this stream-set as processed since we won't need to re-examine it ever again.
        // It is important that we do this first before processing any intermediate streams
        // that may be in this stream-set.
        processedStreamSets.add(streamSet);

        // Set partitions of all intermediate streams in this set (if any)
        int partitions = streamWithSetPartitions.get().getPartitionCount();
        for (StreamEdge streamEdge : streamSet.getStreamEdges()) {
          if (streamEdge.getPartitionCount() == StreamEdge.PARTITIONS_UNKNOWN) {
            streamEdge.setPartitionCount(partitions);
            // Add all unprocessed stream-sets in which this intermediate stream appears
            Collection<StreamSet> streamSetsIncludingIntStream = intermediateStreamToStreamSets.get(streamEdge);
            streamSetsIncludingIntStream.stream()
                .filter(s -> !processedStreamSets.contains(s))
                .forEach(streamSets::add);
          }
        }
      }
    }
  }
Sets partition counts of intermediate streams participating in joins operations.
samza-core/src/main/java/org/apache/samza/execution/IntermediateStreamManager.java