Implement relationship range scan

Author: soerenreichardt

cypher/api/storage-engine-adapter/src/main/java/org/neo4j/gds/compat/AbstractInMemoryRelationshipScanCursor.java

@@ -21,13 +21,18 @@
 
 import org.neo4j.gds.core.cypher.CypherGraphStore;
 import org.neo4j.gds.storageengine.InMemoryRelationshipCursor;
+import org.neo4j.internal.recordstorage.InMemoryRelationshipScan;
 import org.neo4j.storageengine.api.AllRelationshipsScan;
 import org.neo4j.storageengine.api.RelationshipSelection;
 import org.neo4j.storageengine.api.StorageRelationshipScanCursor;
 import org.neo4j.token.TokenHolders;
 
+import static java.lang.Math.min;
+
 public abstract class AbstractInMemoryRelationshipScanCursor extends InMemoryRelationshipCursor implements StorageRelationshipScanCursor {
 
+    private long highMark;
+
     public AbstractInMemoryRelationshipScanCursor(CypherGraphStore graphStore, TokenHolders tokenHolders) {
         super(graphStore, tokenHolders);
     }
@@ -37,30 +42,23 @@ public void scan() {
         reset();
         this.sourceId = 0;
         this.selection = RelationshipSelection.ALL_RELATIONSHIPS;
+        this.highMark = maxRelationshipId;
     }
 
     @Override
     public void single(long reference) {
         reset();
         setId(reference - 1);
+        this.highMark = reference;
         this.selection = RelationshipSelection.ALL_RELATIONSHIPS;
 
-        graphStore.relationshipIds().resolveRelationshipId(reference, (nodeId, offset, context) -> {
-            this.sourceId = nodeId;
-            findContextAndInitializeCursor(context);
-
-            for (long i = 0; i < offset; i++) {
-                next();
-            }
-
-            return null;
-        });
+        initializeForRelationshipReference(reference);
     }
 
     @Override
     public boolean next() {
         if (super.next()) {
-            return true;
+            return getId() <= highMark;
         } else {
             this.sourceId++;
             if (this.sourceId >= graphStore.nodeCount()) {
@@ -73,6 +71,33 @@ public boolean next() {
     }
 
     public boolean scanBatch(AllRelationshipsScan scan, int sizeHint) {
-        throw new UnsupportedOperationException();
+        if (getId() != NO_ID) {
+            reset();
+        }
+
+        highMark = maxRelationshipId;
+        return ((InMemoryRelationshipScan) scan).scanBatch(sizeHint, this);
+    }
+
+    public boolean scanRange(long start, long stop) {
+        reset();
+        this.selection = RelationshipSelection.ALL_RELATIONSHIPS;
+        highMark = min(stop, maxRelationshipId);
+
+        initializeForRelationshipReference(start);
+        return true;
+    }
+
+    private void initializeForRelationshipReference(long reference) {
+        graphStore.relationshipIds().resolveRelationshipId(reference, (nodeId, offset, context) -> {
+            this.sourceId = nodeId;
+            findContextAndInitializeCursor(context);
+
+            for (long i = 0; i < offset; i++) {
+                next();
+            }
+            setId(reference - 1);
+            return null;
+        });
     }
 }

cypher/api/storage-engine-adapter/src/main/java/org/neo4j/gds/storageengine/InMemoryRelationshipCursor.java

@@ -48,6 +48,8 @@ public abstract class InMemoryRelationshipCursor extends RelationshipRecord impl
     private final List<PropertyCursor[]> propertyCursorCache;
     private MutableDoubleList propertyValuesCache;
 
+    protected long maxRelationshipId;
+
     protected long sourceId;
     protected long targetId;
     protected RelationshipSelection selection;
@@ -67,6 +69,7 @@ public InMemoryRelationshipCursor(CypherGraphStore graphStore, TokenHolders toke
         this.propertyCursorCache = new ArrayList<>();
         this.propertyValuesCache = new DoubleArrayList();
 
+        this.maxRelationshipId = 0;
         this.graphStore.relationshipIds().registerUpdateListener(this::newRelationshipIdContextAdded);
     }
 
@@ -163,6 +166,7 @@ private void newRelationshipIdContextAdded(RelationshipIds.RelationshipIdContext
         );
         var newSize = this.propertyCursorCache.stream().mapToInt(cursor -> cursor.length).max().orElse(0);
         this.propertyValuesCache = new DoubleArrayList(new double[newSize]);
+        this.maxRelationshipId += relationshipIdContext.relationshipCount();
     }
 
     private boolean progressToNextContext() {
@@ -191,9 +195,13 @@ private boolean progressToNextContext() {
 
     protected void findContextAndInitializeCursor(RelationshipIds.RelationshipIdContext context) {
         for (int i = 0; i < relationshipIdContexts.size(); i++) {
+            if (i > 0) {
+                relationshipTypeOffset += relationshipIdContexts.get(i - 1).relationshipCount();
+            }
             if (relationshipIdContexts.get(i) == context) {
                 this.relationshipContextIndex = i;
                 initializeCursorForContext(context);
+                break;
             }
         }
     }

cypher/api/storage-engine-adapter/src/main/java/org/neo4j/internal/recordstorage/InMemoryRelationshipScan.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ * This file contains proprietary code that is only available via a commercial license from Neo4j.
+ * For more information, see https://neo4j.com/contact-us/
+ */
+package org.neo4j.internal.recordstorage;
+
+import org.neo4j.gds.compat.AbstractInMemoryRelationshipScanCursor;
+import org.neo4j.storageengine.api.AllRelationshipsScan;
+
+public class InMemoryRelationshipScan extends BaseRecordScan<AbstractInMemoryRelationshipScanCursor> implements AllRelationshipsScan {
+
+    @Override
+    boolean scanRange(AbstractInMemoryRelationshipScanCursor cursor, long start, long stopInclusive) {
+        return cursor.scanRange(start, stopInclusive);
+    }
+
+    @Override
+    public boolean scanBatch(int sizeHint, AbstractInMemoryRelationshipScanCursor cursor) {
+        return super.scanBatch(sizeHint, cursor);
+    }
+}

cypher/cypher-test/build.gradle

@@ -5,6 +5,7 @@ description = 'Neo4j Graph Data Science :: Storage Engine Testing'
 group = 'org.neo4j.gds'
 
 dependencies {
+    testAnnotationProcessor project(':annotations')
     testAnnotationProcessor group: 'org.immutables', name: 'builder',            version: ver.'immutables'
     testAnnotationProcessor group: 'org.immutables', name: 'value',              version: ver.'immutables'
 

cypher/cypher-test/src/test/java/org/neo4j/gds/storageengine/InMemoryRelationshipScanCursorTest.java

@@ -22,6 +22,7 @@
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.PropertyMapping;
 import org.neo4j.gds.StoreLoaderBuilder;
+import org.neo4j.gds.annotation.ValueClass;
 import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.compat.AbstractInMemoryRelationshipScanCursor;
 import org.neo4j.gds.compat.InMemoryPropertySelection;
@@ -33,7 +34,9 @@
 import org.neo4j.gds.junit.annotation.DisableForNeo4jVersion;
 import org.neo4j.values.storable.ValueGroup;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -118,4 +121,34 @@ void shouldSupportPropertiesForSingleRelationship() {
         assertThat(propertyCursor.propertyType()).isEqualTo(ValueGroup.NUMBER);
         assertThat(propertyCursor.propertyKey()).isEqualTo(relationshipScanCursor.tokenHolders.propertyKeyTokens().getIdByName("relProp"));
     }
+
+    @Test
+    void shouldPerformRangeScan() {
+        relationshipScanCursor.scanRange(1, 3);
+
+        var actualRelationships = new HashMap<Long, Relationship>();
+        while(relationshipScanCursor.next()) {
+            actualRelationships.put(
+                relationshipScanCursor.getId(),
+                ImmutableRelationship.of(
+                    relationshipScanCursor.sourceNodeReference(),
+                    relationshipScanCursor.targetNodeReference()
+                )
+            );
+        }
+
+        var expectedRelationships = Map.of(
+            1L, ImmutableRelationship.of(idFunction.of("a"), idFunction.of("c")),
+            2L, ImmutableRelationship.of(idFunction.of("b"), idFunction.of("c")),
+            3L, ImmutableRelationship.of(idFunction.of("a"), idFunction.of("a"))
+        );
+
+        assertThat(actualRelationships).containsExactlyInAnyOrderEntriesOf(expectedRelationships);
+    }
+
+    @ValueClass
+    interface Relationship {
+        long source();
+        long target();
+    }
 }