private static void connectPointwise(
            ExecutionJobVertex jobVertex,
            IntermediateResult result,
            JobVertexInputInfo jobVertexInputInfo) {

        Map<IndexRange, List<Integer>> consumersByPartition = new LinkedHashMap<>();

        for (ExecutionVertexInputInfo executionVertexInputInfo :
                jobVertexInputInfo.getExecutionVertexInputInfos()) {
            int consumerIndex = executionVertexInputInfo.getSubtaskIndex();
            IndexRange range = executionVertexInputInfo.getPartitionIndexRange();
            consumersByPartition.compute(
                    range,
                    (ignore, consumers) -> {
                        if (consumers == null) {
                            consumers = new ArrayList<>();
                        }
                        consumers.add(consumerIndex);
                        return consumers;
                    });
        }

        consumersByPartition.forEach(
                (range, subtasks) -> {
                    List<ExecutionVertex> taskVertices = new ArrayList<>();
                    List<IntermediateResultPartition> partitions = new ArrayList<>();
                    for (int index : subtasks) {
                        taskVertices.add(jobVertex.getTaskVertices()[index]);
                    }
                    for (int i = range.getStartIndex(); i <= range.getEndIndex(); ++i) {
                        partitions.add(result.getPartitions()[i]);
                    }
                    connectInternal(
                            taskVertices,
                            partitions,
                            result.getResultType(),
                            jobVertex.getGraph().getEdgeManager());
                });
    }
@Test
    void testConnectPointwise() throws Exception {
        int upstream = 4;
        int downstream = 4;

        // use dynamic graph to specify the vertex input info
        ExecutionGraph eg = setupExecutionGraph(upstream, downstream, POINTWISE, true);

        // set partition ranges
        List<IndexRange> partitionRanges =
                Arrays.asList(
                        new IndexRange(0, 0),
                        new IndexRange(0, 0),
                        new IndexRange(1, 2),
                        new IndexRange(3, 3));
        List<ExecutionVertexInputInfo> executionVertexInputInfos = new ArrayList<>();
        for (int i = 0; i < downstream; i++) {
            executionVertexInputInfos.add(
                    new ExecutionVertexInputInfo(
                            // the subpartition range will not be used in edge manager, so set (0,
                            // 0)
                            i, partitionRanges.get(i), new IndexRange(0, 0)));
        }
        final JobVertexInputInfo jobVertexInputInfo =
                new JobVertexInputInfo(executionVertexInputInfos);

        final Iterator<ExecutionJobVertex> vertexIterator =
                eg.getVerticesTopologically().iterator();
        final ExecutionJobVertex producer = vertexIterator.next();
        final ExecutionJobVertex consumer = vertexIterator.next();

        // initialize producer and consumer
        eg.initializeJobVertex(producer, 1L, Collections.emptyMap());
        eg.initializeJobVertex(
                consumer,
                1L,
                Collections.singletonMap(
                        producer.getProducedDataSets()[0].getId(), jobVertexInputInfo));

        IntermediateResult result =
                Objects.requireNonNull(eg.getJobVertex(producer.getJobVertexId()))
                        .getProducedDataSets()[0];
        IntermediateResultPartition partition1 = result.getPartitions()[0];
        IntermediateResultPartition partition2 = result.getPartitions()[1];
        IntermediateResultPartition partition3 = result.getPartitions()[2];
        IntermediateResultPartition partition4 = result.getPartitions()[3];

        ExecutionVertex vertex1 = consumer.getTaskVertices()[0];
        ExecutionVertex vertex2 = consumer.getTaskVertices()[1];
        ExecutionVertex vertex3 = consumer.getTaskVertices()[2];
        ExecutionVertex vertex4 = consumer.getTaskVertices()[3];

        // check consumers of the partitions
        ConsumerVertexGroup consumerVertexGroup1 = partition1.getConsumerVertexGroups().get(0);
        ConsumerVertexGroup consumerVertexGroup2 = partition2.getConsumerVertexGroups().get(0);
        ConsumerVertexGroup consumerVertexGroup3 = partition4.getConsumerVertexGroups().get(0);
        assertThat(consumerVertexGroup1)
                .containsExactlyInAnyOrder(vertex1.getID(), vertex2.getID());
        assertThat(consumerVertexGroup2).containsExactlyInAnyOrder(vertex3.getID());
        assertThat(partition3.getConsumerVertexGroups().get(0)).isEqualTo(consumerVertexGroup2);
        assertThat(consumerVertexGroup3).containsExactlyInAnyOrder(vertex4.getID());

        // check inputs of the execution vertices
        ConsumedPartitionGroup consumedPartitionGroup1 = vertex1.getConsumedPartitionGroup(0);
        ConsumedPartitionGroup consumedPartitionGroup2 = vertex3.getConsumedPartitionGroup(0);
        ConsumedPartitionGroup consumedPartitionGroup3 = vertex4.getConsumedPartitionGroup(0);
        assertThat(consumedPartitionGroup1).containsExactlyInAnyOrder(partition1.getPartitionId());
        assertThat(vertex2.getConsumedPartitionGroup(0)).isEqualTo(consumedPartitionGroup1);
        assertThat(consumedPartitionGroup2)
                .containsExactlyInAnyOrder(
                        partition2.getPartitionId(), partition3.getPartitionId());
        assertThat(consumedPartitionGroup3).containsExactlyInAnyOrder(partition4.getPartitionId());

        // check the consumerVertexGroups and consumedPartitionGroups are properly set
        assertThat(consumerVertexGroup1.getConsumedPartitionGroup())
                .isEqualTo(consumedPartitionGroup1);
        assertThat(consumedPartitionGroup1.getConsumerVertexGroup())
                .isEqualTo(consumerVertexGroup1);
        assertThat(consumerVertexGroup2.getConsumedPartitionGroup())
                .isEqualTo(consumedPartitionGroup2);
        assertThat(consumedPartitionGroup2.getConsumerVertexGroup())
                .isEqualTo(consumerVertexGroup2);
        assertThat(consumerVertexGroup3.getConsumedPartitionGroup())
                .isEqualTo(consumedPartitionGroup3);
        assertThat(consumedPartitionGroup3.getConsumerVertexGroup())
                .isEqualTo(consumerVertexGroup3);
    }