org.apache.flink.runtime.operators.sort.PartialOrderPriorityQueue Maven / Gradle / Ivy
/*
* 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.
*/
/**
* This file is based on source code from the Hadoop Project (http://hadoop.apache.org/), licensed by the Apache
* Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE file distributed with this work for
* additional information regarding copyright ownership.
*/
package org.apache.flink.runtime.operators.sort;
import java.util.AbstractQueue;
import java.util.Arrays;
import java.util.Comparator;
import java.util.Iterator;
import java.util.Queue;
/**
* This class implements a priority-queue, which maintains a partial
* ordering of its elements such that the least element can always be found
* in constant time. Put()'s and pop()'s require log(size) time.
*
*/
public class PartialOrderPriorityQueue extends AbstractQueue implements Queue {
/**
* The heap, organized as an array.
*/
private final T[] heap;
/**
* The comparator used to establish the order between the streams.
*/
private final Comparator comparator;
/**
* The maximum size of the heap.
*/
private final int capacity;
/**
* The current number of elements in the queue.
*/
private int size;
@SuppressWarnings("unchecked")
public PartialOrderPriorityQueue(Comparator comparator, int capacity) {
this.comparator = comparator;
this.capacity = capacity + 1;
this.size = 0;
this.heap = (T[]) new Object[this.capacity];
}
/**
* Determines the ordering of objects in this priority queue.
*
* @param a
* The first element.
* @param b
* The second element.
* @return True, if a < b, false otherwise.
*/
private final boolean lessThan(T a, T b) {
return comparator.compare(a, b) < 0;
}
/**
* Returns the remaining capacity of the backing array.
*
* @return The remaining capacity of the backing array.
*/
public int remainingCapacity() {
return capacity - size;
}
/**
* Adds a buffer to a PriorityQueue in log(size) time. If one tries to
* add more objects than maxSize from initialize a RuntimeException
* (ArrayIndexOutOfBound) is thrown.
*/
public final void put(T element) {
size++;
heap[size] = element;
upHeap();
}
/**
* Adds element to the PriorityQueue in log(size) time if either the
* PriorityQueue is not full, or not lessThan(element, top()).
*
* @param element
* The element to insert,
* @return True, if element is added, false otherwise.
*/
public boolean offer(T element) {
if (size < capacity) {
put(element);
return true;
} else if (size > 0 && !lessThan(element, peek())) {
heap[1] = element;
adjustTop();
return true;
} else {
return false;
}
}
/**
* Returns the least element of the PriorityQueue in constant time, but
* does not remove it from the priority queue.
*
* @return The least element.
*/
public final T peek() {
if (size > 0) {
return heap[1];
} else {
return null;
}
}
/**
* Removes and returns the least element of the PriorityQueue in
* log(size) time.
*
* @return The least element.
*/
public final T poll() {
if (size > 0) {
T result = heap[1]; // save first value
heap[1] = heap[size]; // move last to first
heap[size] = null; // permit GC of objects
size--;
downHeap(); // adjust heap
return result;
} else {
return null;
}
}
/**
* Should be called when the Object at top changes values. Still log(n)
* worst case, but it's at least twice as fast to
*
*
* {
* pq.top().change();
* pq.adjustTop();
* }
*
*
* instead of
*
*
* {
* o = pq.pop();
* o.change();
* pq.push(o);
* }
*
*/
public final void adjustTop() {
downHeap();
}
/**
* Returns the number of elements currently stored in the PriorityQueue.
*
* @return The number of elements in the queue.
*/
public final int size() {
return size;
}
/**
* Removes all entries from the PriorityQueue.
*/
public final void clear() {
for (int i = 0; i <= size; i++) {
heap[i] = null;
}
size = 0;
}
private void upHeap() {
int i = size;
T node = heap[i]; // save bottom node
int j = i >>> 1;
while (j > 0 && lessThan(node, heap[j])) {
heap[i] = heap[j]; // shift parents down
i = j;
j = j >>> 1;
}
heap[i] = node; // install saved node
}
private void downHeap() {
int i = 1;
T node = heap[i]; // save top node
int j = i << 1; // find smaller child
int k = j + 1;
if (k <= size && lessThan(heap[k], heap[j])) {
j = k;
}
while (j <= size && lessThan(heap[j], node)) {
heap[i] = heap[j]; // shift up child
i = j;
j = i << 1;
k = j + 1;
if (k <= size && lessThan(heap[k], heap[j])) {
j = k;
}
}
heap[i] = node; // install saved node
}
@Override
public Iterator iterator() {
return Arrays.asList(heap).iterator();
}
}