Github user azagrebin commented on a diff in the pull request:
https://github.com/apache/flink/pull/6228#discussion_r199567240
--- Diff:
flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/KeyGroupPartitionedPriorityQueue.java
---
@@ -0,0 +1,280 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.state.heap;
+
+import org.apache.flink.runtime.state.InternalPriorityQueue;
+import org.apache.flink.runtime.state.KeyExtractorFunction;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.KeyGroupRangeAssignment;
+import org.apache.flink.util.CloseableIterator;
+import org.apache.flink.util.IOUtils;
+
+import javax.annotation.Nonnegative;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+import java.util.Collection;
+import java.util.Comparator;
+
+/**
+ * This implementation of {@link InternalPriorityQueue} is internally
partitioned into sub-queues per key-group and
+ * essentially works as a heap-of-heaps. Instances will have set semantics
for elements if the sub-queues have set
+ * semantics.
+ *
+ * @param <T> the type of elements in the queue.
+ * @param <PQ> type type of sub-queue used for each key-group partition.
+ */
+public class KeyGroupPartitionedPriorityQueue<T, PQ extends
InternalPriorityQueue<T> & HeapPriorityQueueElement>
+ implements InternalPriorityQueue<T> {
+
+ /** A heap of heap sets. Each sub-heap represents the partition for a
key-group.*/
+ @Nonnull
+ private final HeapPriorityQueue<PQ> keyGroupHeap;
+
+ /** All elements from keyGroupHeap, indexed by their key-group id,
relative to firstKeyGroup. */
+ @Nonnull
+ private final PQ[] keyGroupLists;
--- End diff --
I would suggest to rename it to `KeyGroupQueueSets` along with
`getListForElementKeyGroup` and `keyGroupCache`. Probably some left-over
---
