public static TreeMap<Integer, List<BufferIndexAndChannel>>
            getBuffersByConsumptionPriorityInOrder(
                    List<Integer> nextBufferIndexToConsume,
                    TreeMap<Integer, Deque<BufferIndexAndChannel>> subpartitionToAllBuffers,
                    int expectedSize) {
        if (expectedSize <= 0) {
            return new TreeMap<>();
        }

        PriorityQueue<BufferConsumptionPriorityIterator> heap = new PriorityQueue<>();
        subpartitionToAllBuffers.forEach(
                (subpartitionId, buffers) -> {
                    if (!buffers.isEmpty()) {
                        heap.add(
                                new BufferConsumptionPriorityIterator(
                                        buffers, nextBufferIndexToConsume.get(subpartitionId)));
                    }
                });

        TreeMap<Integer, List<BufferIndexAndChannel>> subpartitionToHighPriorityBuffers =
                new TreeMap<>();
        for (int i = 0; i < expectedSize; i++) {
            if (heap.isEmpty()) {
                break;
            }
            BufferConsumptionPriorityIterator bufferConsumptionPriorityIterator = heap.poll();
            BufferIndexAndChannel bufferIndexAndChannel = bufferConsumptionPriorityIterator.next();
            subpartitionToHighPriorityBuffers
                    .computeIfAbsent(bufferIndexAndChannel.getChannel(), k -> new ArrayList<>())
                    .add(bufferIndexAndChannel);
            // if this iterator has next, re-added it.
            if (bufferConsumptionPriorityIterator.hasNext()) {
                heap.add(bufferConsumptionPriorityIterator);
            }
        }
        // treeMap will ensure that the key are sorted by subpartitionId
        // ascending. Within the same subpartition, the larger the bufferIndex,
        // the higher the consumption priority, reserve the value so that buffers are ordered
        // by (subpartitionId, bufferIndex) ascending.
        subpartitionToHighPriorityBuffers.values().forEach(Collections::reverse);
        return subpartitionToHighPriorityBuffers;
    }
@Test
    void testGetBuffersByConsumptionPriorityInOrder() {
        final int subpartition1 = 0;
        final int subpartition2 = 1;

        final int progress1 = 10;
        final int progress2 = 20;

        TreeMap<Integer, Deque<BufferIndexAndChannel>> subpartitionBuffers = new TreeMap<>();
        List<BufferIndexAndChannel> subpartitionBuffers1 =
                createBufferIndexAndChannelsList(
                        subpartition1, progress1, progress1 + 2, progress1 + 6);
        List<BufferIndexAndChannel> subpartitionBuffers2 =
                createBufferIndexAndChannelsList(
                        subpartition2, progress2 + 1, progress2 + 2, progress2 + 5);
        subpartitionBuffers.put(subpartition1, new ArrayDeque<>(subpartitionBuffers1));
        subpartitionBuffers.put(subpartition2, new ArrayDeque<>(subpartitionBuffers2));

        TreeMap<Integer, List<BufferIndexAndChannel>> buffersByConsumptionPriorityInOrder =
                HsSpillingStrategyUtils.getBuffersByConsumptionPriorityInOrder(
                        Arrays.asList(progress1, progress2), subpartitionBuffers, 5);

        assertThat(buffersByConsumptionPriorityInOrder).hasSize(2);
        assertThat(buffersByConsumptionPriorityInOrder.get(subpartition1))
                .isEqualTo(subpartitionBuffers1.subList(1, 3));
        assertThat(buffersByConsumptionPriorityInOrder.get(subpartition2))
                .isEqualTo(subpartitionBuffers2.subList(0, 3));
    }