Fix NodeFilteredGraph::degree

Co-authored-by: Adam Schill Collberg <[email protected]>

Author: FlorentinD

core/src/main/java/org/neo4j/graphalgo/core/huge/NodeFilteredGraph.java

@@ -19,6 +19,7 @@
  */
 package org.neo4j.graphalgo.core.huge;
 
+import org.apache.commons.lang3.mutable.MutableInt;
 import org.neo4j.collection.primitive.PrimitiveLongIterable;
 import org.neo4j.collection.primitive.PrimitiveLongIterator;
 import org.neo4j.graphalgo.api.FilterGraph;
@@ -28,16 +29,32 @@
 import org.neo4j.graphalgo.api.RelationshipIntersect;
 import org.neo4j.graphalgo.api.RelationshipWithPropertyConsumer;
 import org.neo4j.graphalgo.core.loading.IdMap;
+import org.neo4j.graphalgo.core.utils.paged.AllocationTracker;
+import org.neo4j.graphalgo.core.utils.paged.HugeIntArray;
 
 import java.util.Collection;
 import java.util.function.LongPredicate;
 
 public class NodeFilteredGraph extends FilterGraph implements IdMapGraph {
 
+    private static final int NO_DEGREE = -1;
+
     private final IdMap filteredIdMap;
+    private long relationshipCount;
+    private final HugeIntArray degreeCache;
+
+    public NodeFilteredGraph(HugeGraph originalGraph, IdMap filteredIdMap, AllocationTracker allocationTracker) {
+        super(originalGraph);
+        this.relationshipCount = -1;
+        this.filteredIdMap = filteredIdMap;
+        this.degreeCache = HugeIntArray.newArray(filteredIdMap.nodeCount(), allocationTracker);
+        degreeCache.fill(NO_DEGREE);
+    }
+
+    public NodeFilteredGraph(HugeGraph originalGraph, IdMap filteredIdMap, HugeIntArray degreeCache) {
+        super(originalGraph);
 
-    public NodeFilteredGraph(HugeGraph graph, IdMap filteredIdMap) {
-        super(graph);
+        this.degreeCache = degreeCache;
         this.filteredIdMap = filteredIdMap;
     }
 
@@ -68,7 +85,20 @@ public void forEachNode(LongPredicate consumer) {
 
     @Override
     public int degree(long nodeId) {
-        return super.degree(filteredIdMap.toOriginalNodeId(nodeId));
+        int cachedDegree = degreeCache.get(nodeId);
+        if (cachedDegree != NO_DEGREE) {
+            return cachedDegree;
+        }
+
+        MutableInt degree = new MutableInt();
+
+        forEachRelationship(nodeId, (s, t) -> {
+            degree.increment();
+            return true;
+        });
+        degreeCache.set(nodeId, degree.intValue());
+
+        return degree.intValue();
     }
 
     @Override
@@ -131,7 +161,7 @@ public double relationshipProperty(long sourceNodeId, long targetNodeId) {
 
     @Override
     public IdMapGraph concurrentCopy() {
-        return new NodeFilteredGraph((HugeGraph) graph.concurrentCopy(), filteredIdMap);
+        return new NodeFilteredGraph((HugeGraph) graph.concurrentCopy(), filteredIdMap, degreeCache);
     }
 
     @Override

core/src/main/java/org/neo4j/graphalgo/core/loading/GraphStore.java

@@ -472,7 +472,7 @@ protected IdMapGraph createGraph(
                 );
 
                 if (filteredNodes.isPresent()) {
-                    return new NodeFilteredGraph(initialGraph, filteredNodes.get());
+                    return new NodeFilteredGraph(initialGraph, filteredNodes.get(), tracker);
                 } else {
                     return initialGraph;
                 }

core/src/test/java/org/neo4j/graphalgo/core/huge/NodeFilteredGraphTest.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+package org.neo4j.graphalgo.core.huge;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.neo4j.graphalgo.NodeLabel;
+import org.neo4j.graphalgo.RelationshipType;
+import org.neo4j.graphalgo.StoreLoaderBuilder;
+import org.neo4j.graphalgo.TestDatabaseCreator;
+import org.neo4j.graphalgo.api.Graph;
+import org.neo4j.graphalgo.compat.GraphDbApi;
+import org.neo4j.graphalgo.core.loading.GraphStore;
+import org.neo4j.graphalgo.core.loading.NativeFactory;
+
+import java.util.Optional;
+
+import static java.util.Arrays.asList;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.neo4j.graphalgo.QueryRunner.runQuery;
+
+public class NodeFilteredGraphTest {
+
+    private GraphDbApi db;
+
+    private static final String DB_CYPHER =
+        " CREATE" +
+        " (a:Person)," +
+        " (b:Ignore:Person)," +
+        " (c:Ignore:Person)," +
+        " (d:Person)," +
+        " (e:Ignore)," +
+        " (a)-[:REL]->(b)," +
+        " (a)-[:REL]->(e)";
+
+    @BeforeEach
+    void setup() {
+        db = TestDatabaseCreator.createTestDatabase();
+        runQuery(db, DB_CYPHER);
+    }
+
+    @AfterEach
+    void teardown() {
+        db.shutdown();
+    }
+
+
+    @Test
+    void filterDegree() {
+        GraphStore graphStore = new StoreLoaderBuilder()
+            .api(db)
+            .nodeLabels(asList("Person", "Ignore"))
+            .addRelationshipType("REL")
+            .build()
+            .graphStore(NativeFactory.class);
+
+        Graph graph = graphStore.getGraph(
+            asList(NodeLabel.of("Person")),
+            asList(RelationshipType.of("REL")),
+            Optional.empty(),
+            4
+        );
+
+        long nodeIdOfA = graph.toMappedNodeId(0);
+        assertEquals(1L, graph.degree(nodeIdOfA));
+    }
+
+}