Github user StefanRRichter commented on a diff in the pull request: https://github.com/apache/flink/pull/6228#discussion_r199813559 --- Diff: flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapPriorityQueue.java --- @@ -0,0 +1,341 @@ +/* + * 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.util.CloseableIterator; + +import javax.annotation.Nonnegative; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +import java.util.Arrays; +import java.util.Collection; +import java.util.Comparator; +import java.util.Iterator; +import java.util.NoSuchElementException; + +import static org.apache.flink.util.CollectionUtil.MAX_ARRAY_SIZE; + +/** + * Basic heap-based priority queue for {@link HeapPriorityQueueElement} objects. This heap supports fast deletes + * because it manages position indexes of the contained {@link HeapPriorityQueueElement}s. The heap implementation is + * a simple binary tree stored inside an array. Element indexes in the heap array start at 1 instead of 0 to make array + * index computations a bit simpler in the hot methods. Object identification of remove is based on object identity and + * not on equals. + * + * <p>Possible future improvements: + * <ul> + * <li>We could also implement shrinking for the heap.</li> + * </ul> + * + * @param <T> type of the contained elements. + */ +public class HeapPriorityQueue<T extends HeapPriorityQueueElement> implements InternalPriorityQueue<T> { + + /** + * The index of the head element in the array that represents the heap. + */ + private static final int QUEUE_HEAD_INDEX = 1; + + /** + * Comparator for the contained elements. + */ + private final Comparator<T> elementComparator; + + /** + * The array that represents the heap-organized priority queue. + */ + private T[] queue; + + /** + * The current size of the priority queue. + */ + private int size; + + /** + * Creates an empty {@link HeapPriorityQueue} with the requested initial capacity. + * + * @param elementComparator comparator for the contained elements. + * @param minimumCapacity the minimum and initial capacity of this priority queue. + */ + @SuppressWarnings("unchecked") + public HeapPriorityQueue( + @Nonnull Comparator<T> elementComparator, + @Nonnegative int minimumCapacity) { + + this.elementComparator = elementComparator; + this.queue = (T[]) new HeapPriorityQueueElement[QUEUE_HEAD_INDEX + minimumCapacity]; + } + + @Override + @Nullable + public T poll() { + return size() > 0 ? removeElementAtIndex(QUEUE_HEAD_INDEX) : null; + } + + @Override + @Nullable + public T peek() { + return size() > 0 ? queue[QUEUE_HEAD_INDEX] : null; + } + + /** + * Adds the element to add to the heap. This element should not be managed by any other {@link HeapPriorityQueue}. + * + * @return <code>true</code> if the operation changed the head element or if is it unclear if the head element changed. + * Only returns <code>false</code> iff the head element was not changed by this operation. + */ + @Override + public boolean add(@Nonnull T toAdd) { + return addInternal(toAdd); + } + + /** + * This remove is based on object identity, not the result of equals. + * + * @return <code>true</code> if the operation changed the head element or if is it unclear if the head element changed. + * Only returns <code>false</code> iff the head element was not changed by this operation. + */ + @Override + public boolean remove(@Nonnull T toStop) { --- End diff -- ð
---