Termination flag for sampling algorithms - abstract NodeSampler (#9724)

Support termination flag in sampling algorithms

Author: orazve

applications/graph-store-catalog/src/main/java/org/neo4j/gds/applications/graphstorecatalog/DefaultGraphCatalogApplications.java

@@ -963,6 +963,7 @@ public RandomWalkSamplingResult sampleRandomWalkWithRestarts(
         DatabaseId databaseId,
         TaskRegistryFactory taskRegistryFactory,
         UserLogRegistryFactory userLogRegistryFactory,
+        TerminationFlag terminationFlag,
         String graphName,
         String originGraphName,
         Map<String, Object> configuration
@@ -972,6 +973,7 @@ public RandomWalkSamplingResult sampleRandomWalkWithRestarts(
             databaseId,
             taskRegistryFactory,
             userLogRegistryFactory,
+            terminationFlag,
             graphName,
             originGraphName,
             configuration,
@@ -985,6 +987,7 @@ public RandomWalkSamplingResult sampleCommonNeighbourAwareRandomWalk(
         DatabaseId databaseId,
         TaskRegistryFactory taskRegistryFactory,
         UserLogRegistryFactory userLogRegistryFactory,
+        TerminationFlag terminationFlag,
         String graphNameAsString,
         String originGraphName,
         Map<String, Object> configuration
@@ -994,6 +997,7 @@ public RandomWalkSamplingResult sampleCommonNeighbourAwareRandomWalk(
             databaseId,
             taskRegistryFactory,
             userLogRegistryFactory,
+            terminationFlag,
             graphNameAsString,
             originGraphName,
             configuration,
@@ -1104,6 +1108,7 @@ private RandomWalkSamplingResult sampleRandomWalk(
         DatabaseId databaseId,
         TaskRegistryFactory taskRegistryFactory,
         UserLogRegistryFactory userLogRegistryFactory,
+        TerminationFlag terminationFlag,
         String graphNameAsString,
         String originGraphNameAsString,
         Map<String, Object> configuration,
@@ -1125,6 +1130,7 @@ private RandomWalkSamplingResult sampleRandomWalk(
                 userLogRegistryFactory,
                 graphStore,
                 graphProjectConfig,
+                terminationFlag,
                 originGraphName,
                 graphName,
                 configuration,

applications/graph-store-catalog/src/main/java/org/neo4j/gds/applications/graphstorecatalog/GraphCatalogApplications.java

@@ -255,6 +255,7 @@ RandomWalkSamplingResult sampleRandomWalkWithRestarts(
         DatabaseId databaseId,
         TaskRegistryFactory taskRegistryFactory,
         UserLogRegistryFactory userLogRegistryFactory,
+        TerminationFlag terminationFlag,
         String graphName,
         String originGraphName,
         Map<String, Object> configuration
@@ -265,6 +266,7 @@ RandomWalkSamplingResult sampleCommonNeighbourAwareRandomWalk(
         DatabaseId databaseId,
         TaskRegistryFactory taskRegistryFactory,
         UserLogRegistryFactory userLogRegistryFactory,
+        TerminationFlag terminationFlag,
         String graphName,
         String originGraphName,
         Map<String, Object> configuration

applications/graph-store-catalog/src/main/java/org/neo4j/gds/applications/graphstorecatalog/GraphSamplingApplication.java

@@ -32,6 +32,7 @@
 import org.neo4j.gds.graphsampling.GraphSampleConstructor;
 import org.neo4j.gds.graphsampling.RandomWalkSamplerType;
 import org.neo4j.gds.logging.Log;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Map;
 
@@ -50,6 +51,7 @@ RandomWalkSamplingResult sample(
         UserLogRegistryFactory userLogRegistryFactory,
         GraphStore graphStore,
         GraphProjectConfig graphProjectConfig,
+        TerminationFlag terminationFlag,
         GraphName originGraphName,
         GraphName graphName,
         Map<String, Object> configuration,
@@ -73,7 +75,8 @@ RandomWalkSamplingResult sample(
                 samplerConfig,
                 graphStore,
                 samplerAlgorithm,
-                progressTracker
+                progressTracker,
+                terminationFlag
             );
             var sampledGraphStore = graphSampleConstructor.compute();
 

applications/graph-store-catalog/src/test/java/org/neo4j/gds/applications/graphstorecatalog/GraphSamplingApplicationTest.java

@@ -38,6 +38,7 @@
 import org.neo4j.gds.extension.IdFunction;
 import org.neo4j.gds.extension.Inject;
 import org.neo4j.gds.logging.Log;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.List;
 import java.util.Map;
@@ -118,6 +119,7 @@ void shouldSampleRWR(Map<String, Object> mapConfiguration, long expectedNodeCoun
             EmptyUserLogRegistryFactory.INSTANCE,
             graphStore,
             GraphProjectConfig.emptyWithName("user", "graph"),
+            TerminationFlag.RUNNING_TRUE,
             GraphName.parse("graph"),
             GraphName.parse("sample"),
             mapConfiguration,
@@ -157,6 +159,7 @@ void shouldSampleCNARW(Map<String, Object> mapConfiguration, long expectedNodeCo
             EmptyUserLogRegistryFactory.INSTANCE,
             graphStore,
             GraphProjectConfig.emptyWithName("user", "graph"),
+            TerminationFlag.RUNNING_TRUE,
             GraphName.parse("graph"),
             GraphName.parse("sample"),
             mapConfiguration,
@@ -196,6 +199,7 @@ void shouldUseSingleStartNodeRWR(double samplingRatio, long expectedStartNodeCou
             EmptyUserLogRegistryFactory.INSTANCE,
             graphStore,
             GraphProjectConfig.emptyWithName("user", "graph"),
+            TerminationFlag.RUNNING_TRUE,
             GraphName.parse("graph"),
             GraphName.parse("sample"),
             Map.of(
@@ -239,6 +243,7 @@ void shouldUseSingleStartNodeCNARW(double samplingRatio, long expectedStartNodeC
             EmptyUserLogRegistryFactory.INSTANCE,
             graphStore,
             GraphProjectConfig.emptyWithName("user", "graph"),
+            TerminationFlag.RUNNING_TRUE,
             GraphName.parse("graph"),
             GraphName.parse("sample"),
             Map.of(

graph-sampling/src/main/java/org/neo4j/gds/graphsampling/GraphSampleConstructor.java

@@ -45,6 +45,7 @@
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.core.utils.progress.tasks.Tasks;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.List;
 import java.util.Map;
@@ -57,18 +58,21 @@ public class GraphSampleConstructor {
     private final GraphStore inputGraphStore;
     private final NodesSampler nodesSampler;
     private final ProgressTracker progressTracker;
+    private final TerminationFlag terminationFlag;
 
     public GraphSampleConstructor(
         GraphSampleAlgoConfig config,
         GraphStore inputGraphStore,
         NodesSampler nodesSampler,
-        ProgressTracker progressTracker
+        ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
     ) {
         this.config = config;
         this.concurrency = config.concurrency();
         this.inputGraphStore = inputGraphStore;
         this.nodesSampler = nodesSampler;
         this.progressTracker = progressTracker;
+        this.terminationFlag = terminationFlag;
     }
 
     public GraphStore compute() {
@@ -79,6 +83,8 @@ public GraphStore compute() {
             config.internalRelationshipTypes(inputGraphStore),
             config.relationshipWeightProperty()
         );
+        nodesSampler.setTerminationFlag(terminationFlag);
+
         var sampledNodesBitSet = nodesSampler.compute(inputGraph, progressTracker);
 
         progressTracker.beginSubTask("Construct graph");

graph-sampling/src/main/java/org/neo4j/gds/graphsampling/NodesSampler.java

@@ -24,11 +24,25 @@
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
+import org.neo4j.gds.termination.TerminationFlag;
 
-public interface NodesSampler {
-    HugeAtomicBitSet compute(Graph inputGraph, ProgressTracker progressTracker);
+public abstract class NodesSampler {
+    protected abstract HugeAtomicBitSet compute(
+        Graph inputGraph,
+        ProgressTracker progressTracker
+    );
 
-    Task progressTask(GraphStore graphStore);
+    protected abstract Task progressTask(GraphStore graphStore);
 
-    String progressTaskName();
+    protected abstract String progressTaskName();
+
+    protected volatile TerminationFlag terminationFlag = TerminationFlag.RUNNING_TRUE;
+
+    public void setTerminationFlag(TerminationFlag terminationFlag) {
+        this.terminationFlag = terminationFlag;
+    }
+
+    public TerminationFlag getTerminationFlag() {
+        return terminationFlag;
+    }
 }

graph-sampling/src/main/java/org/neo4j/gds/graphsampling/RandomWalkBasedNodesSampler.java

@@ -19,8 +19,8 @@
  */
 package org.neo4j.gds.graphsampling;
 
-public interface RandomWalkBasedNodesSampler extends NodesSampler {
+public abstract class RandomWalkBasedNodesSampler extends NodesSampler {
 
-    long startNodesCount();
+    public abstract long startNodesCount();
 
 }

graph-sampling/src/main/java/org/neo4j/gds/graphsampling/samplers/NodeLabelHistogram.java

@@ -30,6 +30,7 @@
 import org.neo4j.gds.core.concurrency.RunWithConcurrency;
 import org.neo4j.gds.core.utils.partition.PartitionUtils;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Comparator;
 import java.util.Optional;
@@ -45,7 +46,9 @@ interface Result {
         LongLongHashMap histogram();
     }
 
-    public static Result compute(Graph inputGraph, Concurrency concurrency, ProgressTracker progressTracker) {
+    public static Result compute(Graph inputGraph, Concurrency concurrency, ProgressTracker progressTracker,
+        TerminationFlag terminationFlag
+        ) {
         progressTracker.beginSubTask("Count node labels");
         progressTracker.setSteps(inputGraph.nodeCount());
 
@@ -61,6 +64,7 @@ public static Result compute(Graph inputGraph, Concurrency concurrency, Progress
             concurrency,
             inputGraph.nodeCount(),
             partition -> (Runnable) () -> {
+                terminationFlag.assertRunning();
                 var labelCount = new LongLongHashMap();
                 partition.consume(nodeId -> {
                     labelCount.addTo(

graph-sampling/src/main/java/org/neo4j/gds/graphsampling/samplers/SeenNodes.java

@@ -26,6 +26,7 @@
 import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Arrays;
 
@@ -40,14 +41,16 @@ public interface SeenNodes {
     long totalExpectedNodes();
 
     static SeenNodes create(
-        Graph inputGraph, ProgressTracker progressTracker, boolean nodeLabelStratification,
+        Graph inputGraph, ProgressTracker progressTracker, TerminationFlag terminationFlag,
+        boolean nodeLabelStratification,
         Concurrency concurrency, double samplingRatio
     ) {
         if (nodeLabelStratification) {
             var nodeLabelHistogram = NodeLabelHistogram.compute(
                 inputGraph,
                 concurrency,
-                progressTracker
+                progressTracker,
+                terminationFlag
             );
 
             return new SeenNodes.SeenNodesByLabelSet(inputGraph, nodeLabelHistogram, samplingRatio);

graph-sampling/src/main/java/org/neo4j/gds/graphsampling/samplers/rw/Walker.java

@@ -27,6 +27,7 @@
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.graphsampling.samplers.SeenNodes;
 import org.neo4j.gds.graphsampling.samplers.rw.rwr.RandomWalkWithRestarts;
+import org.neo4j.gds.termination.TerminationFlag;
 
 import java.util.Optional;
 import java.util.SplittableRandom;
@@ -41,6 +42,7 @@ public class Walker implements Runnable {
     protected final Graph inputGraph;
     private final double restartProbability;
     protected final ProgressTracker progressTracker;
+    private final TerminationFlag terminationFlag;
 
     private final LongSet startNodesUsed;
 
@@ -55,6 +57,7 @@ public Walker(
         Graph inputGraph,
         double restartProbability,
         ProgressTracker progressTracker,
+        TerminationFlag terminationFlag,
         NextNodeStrategy nextNodeStrategy
     ) {
         this.seenNodes = seenNodes;
@@ -65,6 +68,7 @@ public Walker(
         this.inputGraph = inputGraph;
         this.restartProbability = restartProbability;
         this.progressTracker = progressTracker;
+        this.terminationFlag = terminationFlag;
         this.startNodesUsed = new LongHashSet();
         this.nextNodeStrategy = nextNodeStrategy;
     }
@@ -78,7 +82,7 @@ public void run() {
         int nodesConsidered = 1;
         int walksLeft = (int) Math.round(walkQualities.nodeQuality(currentStartNodePosition) * RandomWalkWithRestarts.MAX_WALKS_PER_START);
 
-        while (!seenNodes.hasSeenEnough()) {
+        while (!seenNodes.hasSeenEnough() && terminationFlag.running()) {
             if (seenNodes.addNode(currentNode)) {
                 addedNodes++;
             }
@@ -118,6 +122,7 @@ public void run() {
                 nodesConsidered++;
             }
         }
+        terminationFlag.assertRunning();
     }
 
     private double computeDegree(long currentNode) {