Add block aligned partitioning

Co-authored-by: Paul Horn <[email protected]>

Author: s1ck

core/src/main/java/org/neo4j/graphalgo/core/utils/partition/PartitionUtils.java

@@ -19,12 +19,16 @@
  */
 package org.neo4j.graphalgo.core.utils.partition;
 
+import com.carrotsearch.hppc.AbstractIterator;
 import org.neo4j.graphalgo.api.Graph;
 import org.neo4j.graphalgo.core.concurrency.ParallelUtil;
 import org.neo4j.graphalgo.core.utils.BitUtil;
 import org.neo4j.graphalgo.core.utils.collection.primitive.PrimitiveLongIterator;
+import org.neo4j.graphalgo.core.utils.paged.HugeCursor;
+import org.neo4j.graphalgo.core.utils.paged.HugeLongArray;
 
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.function.Function;
 
@@ -123,4 +127,80 @@ private static <TASK> List<TASK> tasks(
         }
         return result;
     }
+
+    public static <TASK> Iterator<TASK> blockAlignedPartitioning(
+        HugeLongArray sortedIds,
+        int blockShift,
+        Function<Partition, TASK> taskCreator
+    ) {
+        return new BlockAlignedPartitionIterator<>(sortedIds, blockShift, taskCreator);
+    }
+
+    private static class BlockAlignedPartitionIterator<TASK> extends AbstractIterator<TASK> {
+        private final HugeCursor<long[]> cursor;
+        private final long size;
+        private final int blockShift;
+        private final Function<Partition, TASK> taskCreator;
+
+        private int prevBlockId;
+        private long blockStart;
+        private boolean done;
+        private int lastIndex;
+
+        BlockAlignedPartitionIterator(
+            HugeLongArray sortedIds,
+            int blockShift,
+            Function<Partition, TASK> taskCreator
+        ) {
+            this.size = sortedIds.size();
+            this.blockShift = blockShift;
+            this.taskCreator = taskCreator;
+            this.cursor = sortedIds.initCursor(sortedIds.newCursor());
+            this.prevBlockId = 0;
+            this.blockStart = 0L;
+            this.done = false;
+            this.lastIndex = Integer.MAX_VALUE;
+        }
+
+        @Override
+        protected TASK fetch() {
+            if (this.done) {
+                return done();
+            }
+
+            long base = cursor.base;
+            int limit = cursor.limit;
+            long[] array = cursor.array;
+            int prevBlockId = this.prevBlockId;
+            int blockShift = this.blockShift;
+
+            for (int i = lastIndex; i < limit; i++) {
+                long originalId = array[i];
+                int blockId = (int) (originalId >>> blockShift);
+                if (blockId != prevBlockId) {
+                    long internalId = base + i;
+                    prevBlockId = blockId;
+
+                    if (internalId > 0) {
+                        var partition = ImmutablePartition.of(blockStart, internalId - blockStart);
+                        this.blockStart = internalId;
+                        this.prevBlockId = prevBlockId;
+                        this.lastIndex = i;
+                        return taskCreator.apply(partition);
+                    }
+                }
+            }
+
+            if (cursor.next()) {
+                this.prevBlockId = prevBlockId;
+                this.lastIndex = cursor.offset;
+                return fetch();
+            }
+
+            var partition = ImmutablePartition.of(blockStart, size - blockStart);
+            this.done = true;
+
+            return taskCreator.apply(partition);
+        }
+    }
 }

core/src/test/java/org/neo4j/graphalgo/core/utils/PartitionUtilsTest.java

@@ -29,13 +29,16 @@
 import org.neo4j.graphalgo.beta.generator.RandomGraphGenerator;
 import org.neo4j.graphalgo.beta.generator.RelationshipDistribution;
 import org.neo4j.graphalgo.core.utils.mem.AllocationTracker;
+import org.neo4j.graphalgo.core.utils.paged.HugeLongArray;
+import org.neo4j.graphalgo.core.utils.partition.ImmutablePartition;
 import org.neo4j.graphalgo.core.utils.partition.Partition;
 import org.neo4j.graphalgo.core.utils.partition.PartitionUtils;
 
 import java.util.List;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.within;
@@ -180,6 +183,38 @@ void testDegreePartitioningWithNodeFilter() {
         assertEquals(3, partitions.get(0).nodeCount());
     }
 
+    @Test
+    void testBlockAlignedPartitioning() {
+        var blockShift = 3; // 2^3 = 8 ids per block
+
+        var sortedArray = HugeLongArray.of(
+            /* block 0 */ 0, 3, 5,
+            /* block 1 */ 9, 10, 11, 12, 13, 14, 15,
+            /* block 2 */
+            /* block 3 */ 24, 28, 30, 31,
+            /* block 4 */ 32
+        );
+
+        var partitionIterator = PartitionUtils.blockAlignedPartitioning(
+            sortedArray,
+            blockShift,
+            partition -> partition
+        );
+
+        // 😿 Java
+        var partitions = StreamSupport
+            .stream(((Iterable<Partition>) () -> partitionIterator).spliterator(), false)
+            .collect(Collectors.toList());
+
+        assertThat(partitions)
+            .containsExactly(
+                ImmutablePartition.of(0, 3),
+                ImmutablePartition.of(3, 7),
+                ImmutablePartition.of(10, 4),
+                ImmutablePartition.of(14, 1)
+            );
+    }
+
     static class TestTask implements Runnable {
 
         public final long start;