Refactor scan and base relationship cursor

Author: soerenreichardt

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

@@ -87,17 +87,4 @@ public boolean scanRange(long start, long stop) {
         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

@@ -39,7 +39,11 @@
 import java.util.Arrays;
 import java.util.List;
 
-public abstract class InMemoryRelationshipCursor extends RelationshipRecord implements RelationshipVisitor<RuntimeException>, StorageRelationshipCursor {
+import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
+
+public abstract class InMemoryRelationshipCursor
+    extends RelationshipRecord
+    implements RelationshipVisitor<RuntimeException>, StorageRelationshipCursor, RelationshipIds.UpdateListener {
 
     protected final CypherGraphStore graphStore;
     protected final TokenHolders tokenHolders;
@@ -70,7 +74,7 @@ public InMemoryRelationshipCursor(CypherGraphStore graphStore, TokenHolders toke
         this.propertyValuesCache = new DoubleArrayList();
 
         this.maxRelationshipId = 0;
-        this.graphStore.relationshipIds().registerUpdateListener(this::newRelationshipIdContextAdded);
+        this.graphStore.relationshipIds().registerUpdateListener(this);
     }
 
     @Override
@@ -122,6 +126,21 @@ public boolean next() {
         }
     }
 
+    @Override
+    public void onRelationshipIdsAdded(RelationshipIds.RelationshipIdContext relationshipIdContext) {
+        this.relationshipIdContexts.add(relationshipIdContext);
+        this.adjacencyCursorCache.add(relationshipIdContext.adjacencyList().rawAdjacencyCursor());
+        this.propertyCursorCache.add(
+            Arrays
+                .stream(relationshipIdContext.adjacencyProperties())
+                .map(AdjacencyProperties::rawPropertyCursor)
+                .toArray(PropertyCursor[]::new)
+        );
+        var newSize = this.propertyCursorCache.stream().mapToInt(cursor -> cursor.length).max().orElse(0);
+        this.propertyValuesCache = new DoubleArrayList(new double[newSize]);
+        this.maxRelationshipId += relationshipIdContext.relationshipCount();
+    }
+
     @Override
     public void reset() {
         this.relationshipContextIndex = -1;
@@ -133,7 +152,7 @@ public void reset() {
         setId(NO_ID);
     }
 
-    public void resetCursors() {
+    protected void resetCursors() {
         relationshipContextIndex = -1;
         relationshipTypeOffset = 0;
         adjacencyCursor = null;
@@ -155,20 +174,6 @@ public void properties(StoragePropertyCursor propertyCursor, InMemoryPropertySel
         inMemoryCursor.initRelationshipPropertyCursor(this.sourceId, propertyIds, propertyValuesCache, selection);
     }
 
-    private void newRelationshipIdContextAdded(RelationshipIds.RelationshipIdContext relationshipIdContext) {
-        this.relationshipIdContexts.add(relationshipIdContext);
-        this.adjacencyCursorCache.add(relationshipIdContext.adjacencyList().rawAdjacencyCursor());
-        this.propertyCursorCache.add(
-            Arrays
-                .stream(relationshipIdContext.adjacencyProperties())
-                .map(AdjacencyProperties::rawPropertyCursor)
-                .toArray(PropertyCursor[]::new)
-        );
-        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() {
         relationshipContextIndex++;
 
@@ -193,17 +198,29 @@ private boolean progressToNextContext() {
         return true;
     }
 
-    protected void findContextAndInitializeCursor(RelationshipIds.RelationshipIdContext context) {
-        for (int i = 0; i < relationshipIdContexts.size(); i++) {
-            if (i > 0) {
-                relationshipTypeOffset += relationshipIdContexts.get(i - 1).relationshipCount();
+    protected void initializeForRelationshipReference(long reference) {
+        graphStore.relationshipIds().resolveRelationshipId(reference, (nodeId, offset, context) -> {
+            this.sourceId = nodeId;
+            findContextAndInitializeCursor(context);
+
+            for (long i = 0; i < offset; i++) {
+                next();
             }
-            if (relationshipIdContexts.get(i) == context) {
-                this.relationshipContextIndex = i;
-                initializeCursorForContext(context);
-                break;
+            setId(reference - 1);
+            return null;
+        });
+    }
+
+    private void findContextAndInitializeCursor(RelationshipIds.RelationshipIdContext context) {
+        while (progressToNextContext()) {
+            if (relationshipIdContexts.get(relationshipContextIndex) == context) {
+                return;
             }
         }
+        throw new IllegalStateException(formatWithLocale(
+            "No relationship context for relationship type %s was found",
+            context.relationshipType().name()
+        ));
     }
 
     private void initializeCursorForContext(RelationshipIds.RelationshipIdContext context) {