Merge pull request #8163 from brs96/fix-topkcomputer-concurrency

Fix topKComputer concurrency
neo4j · Sep 25, 2023 · 3daef8f · 3daef8f
2 parents aedfe9c + 73d18a7
commit 3daef8f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/proc/machine-learning/src/main/java/org/neo4j/gds/ml/kge/TopKMapComputer.java b/proc/machine-learning/src/main/java/org/neo4j/gds/ml/kge/TopKMapComputer.java
@@ -100,9 +100,9 @@ public KGEPredictResult compute() {
                     concurrency,
                     terminationFlag,
                     stream -> {
-                        LongLongPredicate isCandidateLinkPredicate = isCandidateLink(concurrentGraph.get());
                         stream.forEach(node1 -> {
                             terminationFlag.assertRunning();
+                            LongLongPredicate isCandidateLinkPredicate = isCandidateLink(concurrentGraph.get());
 
                             LinkScorer linkScorer = threadLocalScorer.get();
                             linkScorer.init(node1);
@@ -138,6 +138,7 @@ private long estimateWorkload() {
     }
 
     private LongLongPredicate isCandidateLink(Graph graph) {
+        //exists O(n)
         return (s, t) -> s != t && !graph.exists(s, t);
     }
 }
diff --git a/proc/machine-learning/src/test/java/org/neo4j/gds/ml/kge/TopKMapComputerTest.java b/proc/machine-learning/src/test/java/org/neo4j/gds/ml/kge/TopKMapComputerTest.java
@@ -73,7 +73,7 @@ void shouldComputeTopKMapTransE() {
         var sourceNodes = create(0, 1, 2);
         var targetNodes = create(3, 4, 5);
         var topK = 1;
-        var concurrency = 1;
+        var concurrency = 4;
 
         var computer = new TopKMapComputer(
             graph,
@@ -111,7 +111,7 @@ void shouldComputeTopKMapDistMult() {
         var sourceNodes = create(0, 1, 2);
         var targetNodes = create(3, 4, 5);
         var topK = 1;
-        var concurrency = 1;
+        var concurrency = 4;
 
         var computer = new TopKMapComputer(
             graph,
@@ -153,7 +153,7 @@ void shouldComputeOverCorrectFiltering() {
         var sourceNodes = create(0, 1, 2);
         var targetNodes = create(0, 1, 2, 3);
         var topK = 10;
-        var concurrency = 1;
+        var concurrency = 4;
 
         var computer = new TopKMapComputer(
             graph,