improve run time for terminal selection

Author: IoannisPanagiotas

algo/src/main/java/org/neo4j/gds/steiner/SteinerBasedDeltaStepping.java

@@ -29,6 +29,7 @@
 import org.neo4j.gds.core.utils.paged.HugeAtomicLongArray;
 import org.neo4j.gds.core.utils.paged.HugeLongArray;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.core.utils.queue.HugeLongPriorityQueue;
 import org.neo4j.gds.paths.ImmutablePathResult;
 import org.neo4j.gds.paths.PathResult;
 import org.neo4j.gds.paths.delta.TentativeDistances;
@@ -39,6 +40,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.LongAdder;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -65,11 +67,9 @@ public final class SteinerBasedDeltaStepping extends Algorithm<DijkstraResult> {
     private final TentativeDistances distances;
     private final ExecutorService executorService;
     private long pathIndex;
-
     private final long numOfTerminals;
     private  final BitSet unvisitedTerminal;
     private final BitSet mergedWithSource;
-
     private final LongAdder metTerminals;
 
     SteinerBasedDeltaStepping(
@@ -143,24 +143,13 @@ private void syncPhase(List<SteinerBasedDeltaTask> tasks,int currentBin, AtomicL
         progressTracker.endSubTask();
     }
 
-    private long nextTerminal(){
-        long index=unvisitedTerminal.nextSetBit(0);
-        long bestTerminal=index;
-        double bestDistance=distances.distance(bestTerminal);
-        index=unvisitedTerminal.nextSetBit(index+1);
-        while (index!=-1){
-            double currentDistance=distances.distance(index);
-            if (currentDistance  < bestDistance){
-                bestTerminal=index;
-                bestDistance=currentDistance;
-            }
-            index=unvisitedTerminal.nextSetBit(index+1);
-        }
-        return bestTerminal;
+    private long nextTerminal(HugeLongPriorityQueue terminalQueue) {
+        return (terminalQueue.isEmpty()) ? -1 : terminalQueue.top();
     }
 
     private boolean updateSteinerTree(long terminalId,AtomicLong frontierIndex,List<PathResult> paths, ImmutablePathResult.Builder pathResultBuilder) {
         //add the new path to the solution
+
         paths.add(pathResult(
             pathResultBuilder,
             pathIndex++,
@@ -182,7 +171,7 @@ private boolean updateSteinerTree(long terminalId,AtomicLong frontierIndex,List<
 
     }
 
-    private long tryToUpdateSteinerTree(long oldBin, long currentBin) {
+    private long tryToUpdateSteinerTree(long oldBin, long currentBin, HugeLongPriorityQueue terminalQueue) {
         boolean shouldComputeClosestTerminal = false;
         //delta-Stepping differs by Dijkstra in that it processes the nodes not one-by-one but in batches
         //whereas in dijkstra once we examine a node, we are certain we have found the shortest path to it,
@@ -198,7 +187,8 @@ private long tryToUpdateSteinerTree(long oldBin, long currentBin) {
             shouldComputeClosestTerminal = true;
         }
         if (shouldComputeClosestTerminal) {
-            long terminalId = nextTerminal();
+            long terminalId = nextTerminal(terminalQueue);
+            if (terminalId == -1) return -1;
             if (distances.distance(terminalId) < currentBin * delta) {
                 return terminalId;
             }
@@ -217,12 +207,14 @@ public DijkstraResult compute() {
         var frontierIndex = new AtomicLong(0);
         var frontierSize = new AtomicLong(1);
 
-        List<PathResult>  paths=new ArrayList<>();
+        List<PathResult> paths = new ArrayList<>();
 
         this.frontier.set(currentBin, startNode);
         mergedWithSource.set(startNode);
         this.distances.set(startNode, -1, 0);
 
+        HugeLongPriorityQueue terminalQueue = HugeLongPriorityQueue.min(unvisitedTerminal.size());
+        var terminalQueueLock = new ReentrantLock();
         var tasks = IntStream
             .range(0, concurrency)
             .mapToObj(i -> new SteinerBasedDeltaTask(
@@ -231,7 +223,10 @@ public DijkstraResult compute() {
                 distances,
                 delta,
                 frontierIndex,
-                mergedWithSource
+                mergedWithSource,
+                terminalQueue,
+                terminalQueueLock,
+                unvisitedTerminal
             ))
             .collect(Collectors.toList());
 
@@ -244,10 +239,11 @@ public DijkstraResult compute() {
             // Find smallest non-empty bin across all tasks
             currentBin = tasks.stream().mapToInt(SteinerBasedDeltaTask::minNonEmptyBin).min().orElseThrow();
 
-            long terminalId = tryToUpdateSteinerTree(oldCurrentBin, currentBin);
+            long terminalId = tryToUpdateSteinerTree(oldCurrentBin, currentBin, terminalQueue);
 
             if (terminalId != -1) { //if we are certain that we have found a shortest path to one of the remaining terminals
                 //we update the solution and merge its path to the root
+                terminalQueue.pop();
                 shouldBreak = updateSteinerTree(terminalId, frontierIndex, paths, pathResultBuilder);
                 currentBin = 0;
             } else { //otherwise proceed as normal, sync the contents of the  bucket for each thread to the global queue.

algo/src/main/java/org/neo4j/gds/steiner/SteinerBasedDeltaTask.java

@@ -25,10 +25,12 @@
 import com.carrotsearch.hppc.procedures.LongProcedure;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.core.utils.paged.HugeLongArray;
+import org.neo4j.gds.core.utils.queue.HugeLongPriorityQueue;
 import org.neo4j.gds.paths.delta.TentativeDistances;
 
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.locks.ReentrantLock;
 
 import static org.neo4j.gds.steiner.SteinerBasedDeltaStepping.BIN_SIZE_THRESHOLD;
 import static org.neo4j.gds.steiner.SteinerBasedDeltaStepping.NO_BIN;
@@ -44,23 +46,32 @@ class SteinerBasedDeltaTask implements Runnable {
     private LongArrayList[] localBins;
     private SteinerBasedDeltaStepping.Phase phase = SteinerBasedDeltaStepping.Phase.RELAX;
     private final BitSet mergedToSource;
+    private final BitSet uninvisitedTerminal;
+    private final HugeLongPriorityQueue terminalQueue;
+    private final ReentrantLock terminalQueueLock;
 
     SteinerBasedDeltaTask(
         Graph graph,
         HugeLongArray frontier,
         TentativeDistances distances,
         double delta,
         AtomicLong frontierIndex,
-        BitSet mergedToSource
+        BitSet mergedToSource,
+        HugeLongPriorityQueue terminalQueue,
+        ReentrantLock terminalQueueLock,
+        BitSet uninvisitedTerminal
     ) {
 
         this.graph = graph;
         this.frontier = frontier;
         this.distances = distances;
         this.delta = delta;
         this.frontierIndex = frontierIndex;
-        this.mergedToSource=mergedToSource;
+        this.mergedToSource = mergedToSource;
         this.localBins = new LongArrayList[0];
+        this.terminalQueue = terminalQueue;
+        this.terminalQueueLock = terminalQueueLock;
+        this.uninvisitedTerminal = uninvisitedTerminal;
     }
 
     @Override
@@ -148,6 +159,15 @@ private void tryToUpdate(long sourceNodeId, long targetNodeId,double weight) {
             // CAX failed, retry
             oldDist = witness;
         }
+        if (uninvisitedTerminal.get(targetNodeId)) {
+
+            terminalQueueLock.lock();
+            if (!terminalQueue.containsElement(targetNodeId) || terminalQueue.cost(targetNodeId) > newDist) {
+                terminalQueue.set(targetNodeId, newDist);
+            }
+            terminalQueueLock.unlock();
+
+        }
     }
 
     private void updateFrontier() {