Add capacity constructor to ShardedByteArrayLongMap

Co-Authored-By: Martin Junghanns <martin.junghanns@neo4j.com>

Author: knutwalker

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

@@ -44,6 +44,10 @@ public static Builder builder(Concurrency concurrency) {
         return new Builder(concurrency);
     }
 
+    public static Builder builder(Concurrency concurrency, long capacity) {
+        return new Builder(concurrency, capacity);
+    }
+
     private ShardedByteArrayLongMap(
         HugeObjectArray<byte[]> internalNodeMapping,
         ObjectLongMap<byte[]>[] originalNodeMappingShards,
@@ -108,14 +112,16 @@ private static <S extends MapShard> ShardedByteArrayLongMap build(
         var mapShards = new ObjectLongMap[shards.length];
 
         // ignoring concurrency limitation 🤷
-        Arrays.parallelSetAll(mapShards, idx -> {
-            var shard = shards[idx];
-            var mapping = shard.intoMapping();
-            mapping.forEachKeyValue((originalId, mappedId) -> {
-                internalNodeMapping.set(mappedId, originalId);
-            });
-            return mapping;
-        });
+        Arrays.parallelSetAll(
+            mapShards, idx -> {
+                var shard = shards[idx];
+                var mapping = shard.intoMapping();
+                mapping.forEachKeyValue((originalId, mappedId) -> {
+                    internalNodeMapping.set(mappedId, originalId);
+                });
+                return mapping;
+            }
+        );
 
         return new ShardedByteArrayLongMap(
             internalNodeMapping,
@@ -148,6 +154,15 @@ public boolean equals(byte[] object1, byte[] object2) {
             this.lockWrapper = new AbstractMultiReaderMutableCollection.LockWrapper(lock);
         }
 
+        MapShard(long capacity) {
+            this.mapping = new ObjectLongHashMapWithHashingStrategy<>(
+                new ArrayHashingStrategy(),
+                Math.toIntExact(capacity)
+            );
+            this.lock = new ReentrantLock();
+            this.lockWrapper = new AbstractMultiReaderMutableCollection.LockWrapper(lock);
+        }
+
         final AbstractMultiReaderMutableCollection.LockWrapper acquireLock() {
             this.lock.lock();
             return this.lockWrapper;
@@ -179,6 +194,16 @@ public static final class Builder {
                 .toArray(Shard[]::new);
         }
 
+        Builder(Concurrency concurrency, long capacity) {
+            this.nodeCount = new AtomicLong();
+            int numberOfShards = numberOfShards(concurrency);
+            this.shardShift = Long.SIZE - Integer.numberOfTrailingZeros(numberOfShards);
+            this.shardMask = numberOfShards - 1;
+            this.shards = IntStream.range(0, numberOfShards)
+                .mapToObj(__ -> new Shard(this.nodeCount, BitUtil.ceilDiv(capacity, numberOfShards)))
+                .toArray(Shard[]::new);
+        }
+
         /**
          * Add a node to the mapping.
          *
@@ -208,6 +233,11 @@ private Shard(AtomicLong nextId) {
                 this.nextId = nextId;
             }
 
+            private Shard(AtomicLong nextId, long capacity) {
+                super(capacity);
+                this.nextId = nextId;
+            }
+
             long addNode(byte[] nodeId) {
                 this.assertIsUnderLock();
                 long mappedId = mapping.getIfAbsent(nodeId, IdMap.NOT_FOUND);