@Override
    public void recordPartitionInfo(int partitionIndex, ResultPartitionBytes partitionBytes) {
        // Once all partitions are finished, we can convert the subpartition bytes to aggregated
        // value to reduce the space usage, because the distribution of source splits does not
        // affect the distribution of data consumed by downstream tasks of ALL_TO_ALL edges(Hashing
        // or Rebalancing, we do not consider rare cases such as custom partitions here).
        if (aggregatedSubpartitionBytes == null) {
            super.recordPartitionInfo(partitionIndex, partitionBytes);

            if (subpartitionBytesByPartitionIndex.size() == numOfPartitions) {
                long[] aggregatedBytes = new long[numOfSubpartitions];
                subpartitionBytesByPartitionIndex
                        .values()
                        .forEach(
                                subpartitionBytes -> {
                                    checkState(subpartitionBytes.length == numOfSubpartitions);
                                    for (int i = 0; i < subpartitionBytes.length; ++i) {
                                        aggregatedBytes[i] += subpartitionBytes[i];
                                    }
                                });
                this.aggregatedSubpartitionBytes =
                        Arrays.stream(aggregatedBytes).boxed().collect(Collectors.toList());
                this.subpartitionBytesByPartitionIndex.clear();
            }
        }
    }
@Test
    void testGetBytesWithPartialPartitionInfos() {
        AllToAllBlockingResultInfo resultInfo =
                new AllToAllBlockingResultInfo(new IntermediateDataSetID(), 2, 2, false);
        resultInfo.recordPartitionInfo(0, new ResultPartitionBytes(new long[] {32L, 64L}));

        assertThatThrownBy(resultInfo::getNumBytesProduced)
                .isInstanceOf(IllegalStateException.class);
        assertThatThrownBy(resultInfo::getAggregatedSubpartitionBytes)
                .isInstanceOf(IllegalStateException.class);
    }