Improve performance of LongByteMappingStore

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<LongWritable, Writable,
   Writable> implements MappingStore<LongWritable, ByteWritable> {
-  /** 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<Long, byte[]> 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();
   }
 }

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)));
+  }
+}