LowCostIdentityHashMap.java
/*
* 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.arrow.memory;
import org.apache.arrow.util.Preconditions;
import org.apache.arrow.util.VisibleForTesting;
import org.checkerframework.checker.initialization.qual.Initialized;
import org.checkerframework.checker.initialization.qual.UnderInitialization;
import org.checkerframework.checker.nullness.qual.Nullable;
/**
* Highly specialized IdentityHashMap that implements only partial Map APIs. It incurs low initial
* cost (just two elements by default). It assumes Value includes the Key -
* Implements @ValueWithKeyIncluded iface that provides "getKey" method.
*
* @param <K> Key type
* @param <V> Value type
*/
public class LowCostIdentityHashMap<K, V extends ValueWithKeyIncluded<K>> {
/*
* The internal data structure to hold values.
*/
private @Nullable Object[] elementData; // elementData[index] = null;
/* Actual number of values. */
private int size;
/*
* maximum number of elements that can be put in this map before having to
* rehash.
*/
private int threshold;
private static final int DEFAULT_MIN_SIZE = 1;
/* Default load factor of 0.75; */
private static final int LOAD_FACTOR = 7500;
/** Creates a Map with default expected maximum size. */
public LowCostIdentityHashMap() {
this(DEFAULT_MIN_SIZE);
}
/**
* Creates a Map with the specified maximum size parameter.
*
* @param maxSize The estimated maximum number of entries that will be put in this map.
*/
public LowCostIdentityHashMap(int maxSize) {
if (maxSize >= 0) {
this.size = 0;
threshold = getThreshold(maxSize);
elementData = newElementArrayUnderInitialized(computeElementArraySize());
} else {
throw new IllegalArgumentException();
}
}
private int getThreshold(@UnderInitialization LowCostIdentityHashMap<K, V> this, int maxSize) {
// assign the threshold to maxSize initially, this will change to a
// higher value if rehashing occurs.
return maxSize > 2 ? maxSize : 2;
}
private int computeElementArraySize(@UnderInitialization LowCostIdentityHashMap<K, V> this) {
int arraySize = (int) (((long) threshold * 10000) / LOAD_FACTOR);
// ensure arraySize is positive, the above cast from long to int type
// leads to overflow and negative arraySize if threshold is too big
return arraySize < 0 ? -arraySize : arraySize;
}
/**
* Create a new element array.
*
* @param s the number of elements
* @return Reference to the element array
*/
private Object[] newElementArrayInitialized(
@Initialized LowCostIdentityHashMap<K, V> this, int s) {
return new Object[s];
}
/**
* Create a new element array.
*
* @param s the number of elements
* @return Reference to the element array
*/
private Object[] newElementArrayUnderInitialized(
@UnderInitialization LowCostIdentityHashMap<K, V> this, int s) {
return new Object[s];
}
/**
* Removes all elements from this map, leaving it empty.
*
* @see #isEmpty()
* @see #size()
*/
public void clear() {
size = 0;
for (int i = 0; i < elementData.length; i++) {
elementData[i] = null;
}
}
/**
* Returns whether this map contains the specified key.
*
* @param key the key to search for.
* @return {@code true} if this map contains the specified key, {@code false} otherwise.
*/
public boolean containsKey(K key) {
Preconditions.checkNotNull(key);
int index = findIndex(key, elementData);
return (elementData[index] == null) ? false : ((V) elementData[index]).getKey() == key;
}
/**
* Returns whether this map contains the specified value.
*
* @param value the value to search for.
* @return {@code true} if this map contains the specified value, {@code false} otherwise.
*/
public boolean containsValue(V value) {
Preconditions.checkNotNull(value);
for (int i = 0; i < elementData.length; i++) {
if (elementData[i] == value) {
return true;
}
}
return false;
}
/**
* Returns the value of the mapping with the specified key.
*
* @param key the key.
* @return the value of the mapping with the specified key.
*/
public @Nullable V get(K key) {
Preconditions.checkNotNull(key);
int index = findIndex(key, elementData);
return (elementData[index] == null)
? null
: (((V) elementData[index]).getKey() == key) ? (V) elementData[index] : null;
}
/**
* Returns the index where the key is found at, or the index of the next empty spot if the key is
* not found in this table.
*/
@VisibleForTesting
int findIndex(@Nullable Object key, @Nullable Object[] array) {
int length = array.length;
int index = getModuloHash(key, length);
int last = (index + length - 1) % length;
while (index != last) {
if ((array[index] == null) || ((V) array[index]).getKey() == key) {
/*
* Found the key, or the next empty spot (which means key is not
* in the table)
*/
break;
}
index = (index + 1) % length;
}
return index;
}
@VisibleForTesting
static int getModuloHash(@Nullable Object key, int length) {
return ((System.identityHashCode(key) & 0x7FFFFFFF) % length);
}
/**
* Maps the specified key to the specified value.
*
* @param value the value.
* @return the value of any previous mapping with the specified key or {@code null} if there was
* no such mapping.
*/
public V put(V value) {
Preconditions.checkNotNull(value);
K key = value.getKey();
Preconditions.checkNotNull(key);
int index = findIndex(key, elementData);
// if the key doesn't exist in the table
if (elementData[index] == null || ((V) elementData[index]).getKey() != key) {
if (++size > threshold) {
rehash();
index = findIndex(key, elementData);
}
// insert the key and assign the value to null initially
elementData[index] = null;
}
// insert value to where it needs to go, return the old value
Object result = elementData[index];
elementData[index] = value;
return (V) result;
}
@VisibleForTesting
void rehash() {
int newlength = elementData.length * 15 / 10;
if (newlength == 0) {
newlength = 1;
}
@Nullable Object[] newData = newElementArrayInitialized(newlength);
for (int i = 0; i < elementData.length; i++) {
Object key = (elementData[i] == null) ? null : ((V) elementData[i]).getKey();
if (key != null) {
// if not empty
int index = findIndex(key, newData);
newData[index] = elementData[i];
}
}
elementData = newData;
computeMaxSize();
}
private void computeMaxSize() {
threshold = (int) ((long) elementData.length * LOAD_FACTOR / 10000);
}
/**
* Removes the mapping with the specified key from this map.
*
* @param key the key of the mapping to remove.
* @return the value of the removed mapping, or {@code null} if no mapping for the specified key
* was found.
*/
public @Nullable V remove(K key) {
Preconditions.checkNotNull(key);
boolean hashedOk;
int index;
int next;
int hash;
Object result;
Object object;
index = next = findIndex(key, elementData);
if (elementData[index] == null || ((V) elementData[index]).getKey() != key) {
return null;
}
// store the value for this key
result = elementData[index];
// clear value to allow movement of the rest of the elements
elementData[index] = null;
size--;
// shift the following elements up if needed
// until we reach an empty spot
int length = elementData.length;
while (true) {
next = (next + 1) % length;
object = elementData[next];
if (object == null) {
break;
}
hash = getModuloHash(((V) object).getKey(), length);
hashedOk = hash > index;
if (next < index) {
hashedOk = hashedOk || (hash <= next);
} else {
hashedOk = hashedOk && (hash <= next);
}
if (!hashedOk) {
elementData[index] = object;
index = next;
elementData[index] = null;
}
}
return (V) result;
}
/**
* Returns whether this Map has no elements.
*
* @return {@code true} if this Map has no elements, {@code false} otherwise.
* @see #size()
*/
public boolean isEmpty() {
return size == 0;
}
/**
* Returns the number of mappings in this Map.
*
* @return the number of mappings in this Map.
*/
public int size() {
return size;
}
/**
* Special API to return next value - substitute of regular Map.values.iterator().next().
*
* @return next available value or null if none available
*/
public @Nullable V getNextValue() {
for (int i = 0; i < elementData.length; i++) {
if (elementData[i] != null) {
return (V) elementData[i];
}
}
return null;
}
}