org.apache.flink.runtime.state.heap.AbstractHeapPriorityQueue Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of flink-runtime_2.11 Show documentation
There is a newer version: 1.13.6
/*
 * 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.Iterator;
import java.util.NoSuchElementException;

import static org.apache.flink.util.CollectionUtil.MAX_ARRAY_SIZE;

/**
 * Abstract base class for heap (object array) based implementations of priority queues, with support for fast deletes
 * via {@link HeapPriorityQueueElement}.
 *
 * @param  type of the elements contained in the priority queue.
 */
public abstract class AbstractHeapPriorityQueue
	implements InternalPriorityQueue {

	/** The array that represents the heap-organized priority queue. */
	@Nonnull
	protected T[] queue;

	/** The current size of the priority queue. */
	@Nonnegative
	protected int size;

	@SuppressWarnings("unchecked")
	public AbstractHeapPriorityQueue(@Nonnegative int minimumCapacity) {
		this.queue = (T[]) new HeapPriorityQueueElement[getHeadElementIndex() + minimumCapacity];
		this.size = 0;
	}

	@Override
	@Nullable
	public T poll() {
		return size() > 0 ? removeInternal(getHeadElementIndex()) : null;
	}

	@Override
	@Nullable
	public T peek() {
		// References to removed elements are expected to become set to null.
		return queue[getHeadElementIndex()];
	}

	@Override
	public boolean add(@Nonnull T toAdd) {
		addInternal(toAdd);
		return toAdd.getInternalIndex() == getHeadElementIndex();
	}

	@Override
	public boolean remove(@Nonnull T toRemove) {
		final int elementIndex = toRemove.getInternalIndex();
		removeInternal(elementIndex);
		return elementIndex == getHeadElementIndex();
	}

	@Override
	public boolean isEmpty() {
		return size() == 0;
	}

	@Override
	public int size() {
		return size;
	}

	@Override
	public void addAll(@Nullable Collection toAdd) {

		if (toAdd == null) {
			return;
		}

		resizeForBulkLoad(toAdd.size());

		for (T element : toAdd) {
			add(element);
		}
	}

	@SuppressWarnings({"unchecked"})
	@Nonnull
	public  O[] toArray(O[] out) {
		final int heapArrayOffset = getHeadElementIndex();
		if (out.length < size) {
			return (O[]) Arrays.copyOfRange(queue, heapArrayOffset, heapArrayOffset + size, out.getClass());
		} else {
			System.arraycopy(queue, heapArrayOffset, out, 0, size);
			if (out.length > size) {
				out[size] = null;
			}
			return out;
		}
	}

	/**
	 * Returns an iterator over the elements in this queue. The iterator
	 * does not return the elements in any particular order.
	 *
	 * @return an iterator over the elements in this queue.
	 */
	@Nonnull
	@Override
	public CloseableIterator iterator() {
		return new HeapIterator();
	}

	/**
	 * Clears the queue.
	 */
	public void clear() {
		final int arrayOffset = getHeadElementIndex();
		Arrays.fill(queue, arrayOffset, arrayOffset + size, null);
		size = 0;
	}

	protected void resizeForBulkLoad(int totalSize) {
		if (totalSize > queue.length) {
			int desiredSize = totalSize + (totalSize >>> 3);
			resizeQueueArray(desiredSize, totalSize);
		}
	}

	protected void resizeQueueArray(int desiredSize, int minRequiredSize) {
		if (isValidArraySize(desiredSize)) {
			queue = Arrays.copyOf(queue, desiredSize);
		} else if (isValidArraySize(minRequiredSize)) {
			queue = Arrays.copyOf(queue, MAX_ARRAY_SIZE);
		} else {
			throw new OutOfMemoryError("Required minimum heap size " + minRequiredSize +
				" exceeds maximum size of " + MAX_ARRAY_SIZE + ".");
		}
	}

	protected void moveElementToIdx(T element, int idx) {
		queue[idx] = element;
		element.setInternalIndex(idx);
	}

	/**
	 * Implements how to remove the element at the given index from the queue.
	 *
	 * @param elementIndex the index to remove.
	 * @return the removed element.
	 */
	protected abstract T removeInternal(@Nonnegative int elementIndex);

	/**
	 * Implements how to add an element to the queue.
	 *
	 * @param toAdd the element to add.
	 */
	protected abstract void addInternal(@Nonnull T toAdd);

	/**
	 * Returns the start index of the queue elements in the array.
	 */
	protected abstract int getHeadElementIndex();

	private static boolean isValidArraySize(int size) {
		return size >= 0 && size <= MAX_ARRAY_SIZE;
	}

	/**
	 * {@link Iterator} implementation for {@link HeapPriorityQueue}.
	 * {@link Iterator#remove()} is not supported.
	 */
	private final class HeapIterator implements CloseableIterator {

		private int runningIdx;
		private final int endIdx;

		HeapIterator() {
			this.runningIdx = getHeadElementIndex();
			this.endIdx = runningIdx + size;
		}

		@Override
		public boolean hasNext() {
			return runningIdx < endIdx;
		}

		@Override
		public T next() {
			if (runningIdx >= endIdx) {
				throw new NoSuchElementException("Iterator has no next element.");
			}
			return queue[runningIdx++];
		}

		@Override
		public void close() {
		}
	}
}