diff --git a/giraph-core/src/main/java/org/apache/giraph/mapping/LongByteMappingStore.java b/giraph-core/src/main/java/org/apache/giraph/mapping/LongByteMappingStore.java index ab35cdc6f..fb429e447 100644 --- a/giraph-core/src/main/java/org/apache/giraph/mapping/LongByteMappingStore.java +++ b/giraph-core/src/main/java/org/apache/giraph/mapping/LongByteMappingStore.java @@ -22,7 +22,6 @@ import java.util.Arrays; import java.util.concurrent.ConcurrentMap; -import java.util.concurrent.atomic.AtomicLong; import javax.annotation.concurrent.ThreadSafe; @@ -31,7 +30,6 @@ import org.apache.hadoop.io.ByteWritable; import org.apache.hadoop.io.LongWritable; import org.apache.hadoop.io.Writable; -import org.apache.log4j.Logger; import com.google.common.collect.MapMaker; @@ -48,12 +46,6 @@ public class LongByteMappingStore extends DefaultImmutableClassesGiraphConfigurable implements MappingStore { - /** Logger instance */ - private static final Logger LOG = Logger.getLogger( - LongByteMappingStore.class); - - /** Counts number of entries added */ - private final AtomicLong numEntries = new AtomicLong(0); /** Id prefix to bytesArray index mapping */ private ConcurrentMap concurrentIdToBytes; @@ -84,6 +76,7 @@ public void initialize() { .concurrencyLevel(getConf().getNumInputSplitsThreads()) .makeMap(); idToBytes = new Long2ObjectOpenHashMap<>(upper); + idToBytes.defaultReturnValue(null); } /** @@ -94,11 +87,12 @@ public void initialize() { */ public byte getByteTarget(LongWritable vertexId) { long key = vertexId.get() >>> lowerOrder; - int suffix = (int) (vertexId.get() & lowerBitMask); - if (!idToBytes.containsKey(key)) { + byte[] bs = idToBytes.get(key); + if (bs == null) { return -1; } - return idToBytes.get(key)[suffix]; + int suffix = (int) (vertexId.get() & lowerBitMask); + return bs[suffix]; } @Override @@ -114,7 +108,6 @@ public void addEntry(LongWritable vertexId, ByteWritable target) { } } bytes[(int) (vertexId.get() & lowerBitMask)] = target.get(); - numEntries.getAndIncrement(); // increment count } @Override @@ -138,8 +131,15 @@ public void postFilling() { concurrentIdToBytes = null; } + /** + * Returns the number of entries in the mapping store. This is updated only + * after the mapping has finished loading after {@link #postFilling()} has + * been called. + * + * @return + */ @Override public long getStats() { - return numEntries.longValue(); + return idToBytes.size(); } } diff --git a/giraph-core/src/test/java/org/apache/giraph/mapping/LongByteMappingStoreTest.java b/giraph-core/src/test/java/org/apache/giraph/mapping/LongByteMappingStoreTest.java new file mode 100644 index 000000000..a23da03ff --- /dev/null +++ b/giraph-core/src/test/java/org/apache/giraph/mapping/LongByteMappingStoreTest.java @@ -0,0 +1,39 @@ +package org.apache.giraph.mapping; + +import org.apache.giraph.conf.GiraphConfiguration; +import org.apache.giraph.conf.GiraphConstants; +import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration; +import org.apache.hadoop.io.ByteWritable; +import org.apache.hadoop.io.LongWritable; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +public class LongByteMappingStoreTest { + + @Test + public void test() { + ImmutableClassesGiraphConfiguration conf = + new ImmutableClassesGiraphConfiguration(new GiraphConfiguration()); + GiraphConstants.LB_MAPPINGSTORE_UPPER.setIfUnset(conf, 100); + GiraphConstants.LB_MAPPINGSTORE_LOWER.setIfUnset(conf, 4); + LongByteMappingStore store = new LongByteMappingStore(); + store.setConf(conf); + store.initialize(); + store.addEntry(new LongWritable(1), new ByteWritable((byte) 1)); + store.addEntry(new LongWritable(2), new ByteWritable((byte) 2)); + store.postFilling(); + + assertEquals((byte) 1, store.getByteTarget(new LongWritable(1))); + assertEquals((byte) 2, store.getByteTarget(new LongWritable(2))); + assertEquals((byte) -1, store.getByteTarget(new LongWritable(999))); + + assertEquals(new ByteWritable((byte) 1), + store.getTarget(new LongWritable(1), new ByteWritable((byte) 777))); + assertEquals(new ByteWritable((byte) 2), + store.getTarget(new LongWritable(2), new ByteWritable((byte) 888))); + assertNull(store.getTarget(new LongWritable(3), + new ByteWritable((byte) 555))); + } +}