Use custom hashing strategy for map shards

Author: s1ck

core/src/main/java/org/neo4j/gds/core/utils/paged/ShardedByteArrayLongMap.java

@@ -19,10 +19,11 @@
  */
 package org.neo4j.gds.core.utils.paged;
 
+import org.eclipse.collections.api.block.HashingStrategy;
 import org.eclipse.collections.api.map.primitive.MutableObjectLongMap;
 import org.eclipse.collections.api.map.primitive.ObjectLongMap;
 import org.eclipse.collections.impl.collection.mutable.AbstractMultiReaderMutableCollection;
-import org.eclipse.collections.impl.factory.primitive.ObjectLongMaps;
+import org.eclipse.collections.impl.map.mutable.primitive.ObjectLongHashMapWithHashingStrategy;
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.collections.ha.HugeObjectArray;
 import org.neo4j.gds.core.concurrency.Concurrency;
@@ -125,12 +126,24 @@ private static <S extends MapShard> ShardedByteArrayLongMap build(
 
     abstract static class MapShard {
 
+        private static class ArrayHashingStrategy implements HashingStrategy<byte[]> {
+            @Override
+            public int computeHashCode(byte[] object) {
+                return Arrays.hashCode(object);
+            }
+
+            @Override
+            public boolean equals(byte[] object1, byte[] object2) {
+                return Arrays.equals(object1, object2);
+            }
+        }
+
         private final ReentrantLock lock;
         private final AbstractMultiReaderMutableCollection.LockWrapper lockWrapper;
         final MutableObjectLongMap<byte[]> mapping;
 
         MapShard() {
-            this.mapping = ObjectLongMaps.mutable.empty();
+            this.mapping = new ObjectLongHashMapWithHashingStrategy<>(new ArrayHashingStrategy());
             this.lock = new ReentrantLock();
             this.lockWrapper = new AbstractMultiReaderMutableCollection.LockWrapper(lock);
         }

core/src/test/java/org/neo4j/gds/core/utils/paged/ShardedByteArrayLongMapTest.java

@@ -34,6 +34,7 @@
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.Optional;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -43,7 +44,9 @@ class ShardedByteArrayLongMapTest {
     @Provide
     Arbitrary<byte[][]> nodes() {
         var idGen = Arbitraries.bytes().array(byte[].class).ofSize(10);
-        return Arbitraries.create(idGen::sample).array(byte[][].class);
+        return Arbitraries
+            .create(idGen::sample)
+            .array(byte[][].class);
     }
 
     @Test
@@ -61,8 +64,7 @@ void addSingleNode() {
     void addNodes(@ForAll("nodes") @Size(100) byte[][] nodes) {
         var builder = ShardedByteArrayLongMap.builder(new Concurrency(1));
         for (byte[] node : nodes) {
-            long mapped = builder.addNode(node);
-            assertThat(mapped).isGreaterThanOrEqualTo(0);
+            builder.addNode(node);
         }
         var map = builder.build();
 
@@ -72,6 +74,34 @@ void addNodes(@ForAll("nodes") @Size(100) byte[][] nodes) {
         }
     }
 
+    @Property
+    void addNodesDifferentObject(@ForAll("nodes") @Size(100) byte[][] nodes) {
+        var builder = ShardedByteArrayLongMap.builder(new Concurrency(1));
+        for (byte[] node : nodes) {
+            builder.addNode(node);
+        }
+        var map = builder.build();
+
+        assertThat(map.size()).isEqualTo(nodes.length);
+        for (byte[] node : nodes) {
+            // Ensure that hashCode and equals work correctly for byte arrays
+            // with same elements, but different objects.
+            var nodeCopy = Arrays.copyOf(node, node.length);
+            assertThat(map.toOriginalNodeId(map.toMappedNodeId(nodeCopy))).isEqualTo(node);
+        }
+    }
+
+    @Test
+    void addExistingNode() {
+        byte[] node1 = "foobar".getBytes(StandardCharsets.UTF_8);
+        byte[] node2 = "foobar".getBytes(StandardCharsets.UTF_8);
+        var builder = ShardedByteArrayLongMap.builder(new Concurrency(1));
+        long mappedNode1 = builder.addNode(node1);
+        assertThat(mappedNode1).isGreaterThanOrEqualTo(0);
+        long mappedNode2 = builder.addNode(node2);
+        assertThat(mappedNode2).isEqualTo(-(mappedNode1 + 1));
+    }
+
     @ParameterizedTest
     @ValueSource(ints = {0, 1024, 4096, 5000, 9999})
     void size(int expectedSize) {