Refactor Yens

Co-authored-by: Veselin Nikolov <[email protected]>

Author: IoannisPanagiotas

algo/src/main/java/org/neo4j/gds/paths/yens/Yens.java

@@ -39,6 +39,7 @@
 import java.util.Comparator;
 import java.util.Optional;
 import java.util.PriorityQueue;
+import java.util.function.ToLongBiFunction;
 import java.util.stream.Stream;
 
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
@@ -50,9 +51,10 @@ public final class Yens extends Algorithm<DijkstraResult> {
     private final Graph graph;
     private final ShortestPathYensBaseConfig config;
     private final Dijkstra dijkstra;
-
-    private final LongScatterSet nodeBlackList;
-    private final LongObjectScatterMap<LongHashSet> relationshipBlackList;
+    private final LongScatterSet nodeAvoidList;
+    private final LongObjectScatterMap<LongHashSet> relationshipAvoidList;
+    private final ToLongBiFunction
+        <MutablePathResult, Integer> relationshipAvoidMapper;
 
     /**
      * Configure Yens to compute at most one source-target shortest path.
@@ -63,22 +65,15 @@ public static Yens sourceTarget(
         ProgressTracker progressTracker
     ) {
         // If the input graph is a multi-graph, we need to track
-        // parallel relationships. This is necessary since shortest
+        // parallel relationships ids. This is necessary since shortest
         // paths can visit the same nodes via different relationships.
+        //If not, we need to track which is the next neighbor.
 
-        System.out.println(graph.schema().relationshipSchema().toMap());
-        graph.forEachNode(nodeId -> {
-            graph.forEachRelationship(nodeId, 1.0, (s, t, w) -> {
-                System.out.println(s + "-[" + w + "]->" + t);
-                return true;
-            });
-            return true;
-        });
-
+        boolean shouldTrackRelationships = graph.isMultiGraph();
         var newConfig = ImmutableShortestPathYensBaseConfig
             .builder()
             .from(config)
-            .trackRelationships(graph.isMultiGraph())
+            .trackRelationships(shouldTrackRelationships)
             .build();
         // Init dijkstra algorithm for computing shortest paths
         var dijkstra = Dijkstra.sourceTarget(graph, newConfig, Optional.empty(), progressTracker);
@@ -105,16 +100,35 @@ private Yens(Graph graph, Dijkstra dijkstra, ShortestPathYensBaseConfig config,
         this.config = config;
         // Track nodes and relationships that are skipped in a single iteration.
         // The content of these data structures is reset after each of k iterations.
-        this.nodeBlackList = new LongScatterSet();
-        this.relationshipBlackList = new LongObjectScatterMap<>();
-        // set filter in Dijkstra to respect our blacklists
+        this.nodeAvoidList = new LongScatterSet();
+        this.relationshipAvoidList = new LongObjectScatterMap<>();
+        // set filter in Dijkstra to respect our list of relationships to avoid
         this.dijkstra = dijkstra;
+
+        if (config.trackRelationships()) {
+            // if we are in a multi-graph, we  must store the relationships ids as they are
+            //since two nodes may be connected by multiple relationships and we must know which to avoid
+            relationshipAvoidMapper = (path, position) -> path.relationship(position);
+        } else {
+            //otherwise the graph has surely no parallel edges, we do not need to explicitly store relationship ids
+            //we can just store endpoints, so that we know which nodes a node should avoid
+            relationshipAvoidMapper = (path, position) -> path.node(position + 1);
+        }
         dijkstra.withRelationshipFilter((source, target, relationshipId) ->
-            !nodeBlackList.contains(target) &&
-            !(relationshipBlackList.getOrDefault(source, EMPTY_SET).contains(relationshipId)) &&
-            !(relationshipBlackList.getOrDefault(source, EMPTY_SET).contains(-target - 1)));
+            !nodeAvoidList.contains(target)
+            && !shouldAvoidRelationship(source, target, relationshipId)
+
+        );
     }
 
+    private boolean shouldAvoidRelationship(long source, long target, long relationshipId) {
+        long forbidden = target;
+        if (config.trackRelationships()) {
+            forbidden = relationshipId;
+        }
+        return relationshipAvoidList.getOrDefault(source, EMPTY_SET).contains(forbidden);
+
+    }
 
     @Override
     public DijkstraResult compute() {
@@ -150,22 +164,21 @@ public DijkstraResult compute() {
                     // Filter relationships that are part of the previous
                     // shortest paths which share the same root path.
                     if (rootPath.matchesExactly(path, n + 1)) {
-                        System.out.println(i + ": " + rootPath + " |" + prevPath);
-                        var relationshipId = graph.isMultiGraph() ? path.relationship(n) : -(1 + path.node(n + 1));
+                        var relationshipId = relationshipAvoidMapper.applyAsLong(path, n);
 
-                        var neighbors = relationshipBlackList.get(spurNode);
+                        var neighbors = relationshipAvoidList.get(spurNode);
 
                         if (neighbors == null) {
                             neighbors = new LongHashSet();
-                            relationshipBlackList.put(spurNode, neighbors);
+                            relationshipAvoidList.put(spurNode, neighbors);
                         }
                         neighbors.add(relationshipId);
                     }
                 }
 
                 // Filter nodes from root path to avoid cyclic path searches.
                 for (int j = 0; j < n; j++) {
-                    nodeBlackList.add(rootPath.node(j));
+                    nodeAvoidList.add(rootPath.node(j));
                 }
 
                 // Calculate the spur path from the spur node to the sink.
@@ -174,8 +187,8 @@ public DijkstraResult compute() {
                 var spurPath = computeDijkstra(graph.toOriginalNodeId(spurNode));
 
                 // Clear filters for next spur node
-                nodeBlackList.clear();
-                relationshipBlackList.clear();
+                nodeAvoidList.clear();
+                relationshipAvoidList.clear();
 
                 // No new candidate from this spur node, continue with next node.
                 if (spurPath.isEmpty()) {
@@ -201,10 +214,7 @@ public DijkstraResult compute() {
         progressTracker.endSubTask();
 
         progressTracker.endSubTask();
-        System.out.println("----");
-        for (var path : kShortestPaths) {
-            System.out.println(path);
-        }
+    
         return new DijkstraResult(kShortestPaths.stream().map(MutablePathResult::toPathResult));
     }
 

algo/src/test/java/org/neo4j/gds/paths/yens/YensTest.java

@@ -168,6 +168,7 @@ static Stream<List<String>> pathInput() {
     void compute(Collection<String> expectedPaths) {
 
         assertResult(graph, idFunction, expectedPaths, false);
+
     }
 
     @Test