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 testGetBuffersByConsumptionPriorityInOrderEmptyExpectedSize() {
        TreeMap<Integer, Deque<BufferIndexAndChannel>> subpartitionToAllBuffers = new TreeMap<>();
        subpartitionToAllBuffers.put(0, createBufferIndexAndChannelsDeque(0, 0, 1));
        subpartitionToAllBuffers.put(1, createBufferIndexAndChannelsDeque(1, 2, 4));
        TreeMap<Integer, List<BufferIndexAndChannel>> buffersByConsumptionPriorityInOrder =
                HsSpillingStrategyUtils.getBuffersByConsumptionPriorityInOrder(
                        Arrays.asList(0, 1), subpartitionToAllBuffers, 0);
        assertThat(buffersByConsumptionPriorityInOrder).isEmpty();
    }