Handle tiny graphs for sampling (#9991)

Author: FlorentinD

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

@@ -58,7 +58,7 @@ static SeenNodes create(
 
         return new SeenNodes.GlobalSeenNodes(
             HugeAtomicBitSet.create(inputGraph.nodeCount()),
-            Math.round(inputGraph.nodeCount() * samplingRatio)
+            Math.max(1, Math.round(inputGraph.nodeCount() * samplingRatio))
         );
     }
 

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

@@ -77,6 +77,12 @@ public HugeAtomicBitSet compute(Graph inputGraph, ProgressTracker progressTracke
             config.samplingRatio()
         );
 
+        if (seenNodes.totalExpectedNodes() == 0) {
+            progressTracker.endSubTask("Sample nodes");
+            return seenNodes.sampledNodes();
+        }
+
+
         progressTracker.beginSubTask("Do common neighbour aware random walks");
         progressTracker.setSteps(seenNodes.totalExpectedNodes());
 

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

@@ -75,6 +75,11 @@ public HugeAtomicBitSet compute(Graph inputGraph, ProgressTracker progressTracke
             config.samplingRatio()
         );
 
+        if (seenNodes.totalExpectedNodes() == 0) {
+            progressTracker.endSubTask("Sample nodes");
+            return seenNodes.sampledNodes();
+        }
+
         progressTracker.beginSubTask("Do random walks");
         progressTracker.setSteps(seenNodes.totalExpectedNodes());
 
@@ -116,16 +121,17 @@ public HugeAtomicBitSet compute(Graph inputGraph, ProgressTracker progressTracke
 
     @Override
     public Task progressTask(GraphStore graphStore) {
+        long sampledNodes = 10 * Math.round(graphStore.nodeCount() * config.samplingRatio());
         if (config.nodeLabelStratification()) {
             return Tasks.task(
                 "Sample nodes",
                 Tasks.leaf("Count node labels", graphStore.nodeCount()),
-                Tasks.leaf("Do random walks", 10 * Math.round(graphStore.nodeCount() * config.samplingRatio()))
+                Tasks.leaf("Do random walks", sampledNodes)
             );
         } else {
             return Tasks.task(
                 "Sample nodes",
-                Tasks.leaf("Do random walks", 10 * Math.round(graphStore.nodeCount() * config.samplingRatio()))
+                Tasks.leaf("Do random walks", sampledNodes)
             );
         }
     }

graph-sampling/src/test/java/org/neo4j/gds/graphsampling/samplers/rw/cnarw/CommonNeighbourAwareRandomWalkTest.java

@@ -22,20 +22,27 @@
 import org.assertj.core.data.Offset;
 import org.assertj.core.data.Percentage;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.neo4j.gds.NodeLabel;
+import org.neo4j.gds.TestProgressTracker;
+import org.neo4j.gds.TestTaskStore;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.core.GraphDimensions;
 import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.utils.paged.HugeAtomicBitSet;
+import org.neo4j.gds.core.utils.progress.LocalTaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.extension.GdlExtension;
 import org.neo4j.gds.extension.GdlGraph;
 import org.neo4j.gds.extension.IdFunction;
 import org.neo4j.gds.extension.Inject;
 import org.neo4j.gds.extension.TestGraph;
 import org.neo4j.gds.graphsampling.config.CommonNeighbourAwareRandomWalkConfig;
 import org.neo4j.gds.graphsampling.config.CommonNeighbourAwareRandomWalkConfigImpl;
+import org.neo4j.gds.logging.GdsTestLog;
 import org.neo4j.gds.mem.MemoryRange;
 import org.neo4j.gds.termination.TerminatedException;
 import org.neo4j.gds.termination.TerminationFlag;
@@ -172,6 +179,15 @@ class CommonNeighbourAwareRandomWalkTest {
     @Inject
     private TestGraph naturalUnionGraph;
 
+    @GdlGraph(graphNamePrefix = "tiny")
+    private static String TINY_GRAPH = "()-->()";
+
+    @Inject
+    private Graph tinyGraph;
+
+    @Inject
+    private GraphStore tinyGraphStore;
+
 
     private Graph getGraph(CommonNeighbourAwareRandomWalkConfig config) {
         return graphStore.getGraph(
@@ -670,4 +686,23 @@ void checkTerminationFlag() {
             .isInstanceOf(TerminatedException.class)
             .hasMessageContaining("The execution has been terminated.");
     }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    void progressTrackingTinyGraph(boolean nodeLabelStratification) {
+        var config = CommonNeighbourAwareRandomWalkConfigImpl.builder()
+            .nodeLabelStratification(nodeLabelStratification)
+            .build();
+
+        var cnar = new CommonNeighbourAwareRandomWalk(config);
+        Task task = cnar.progressTask(tinyGraphStore);
+
+        TestTaskStore taskStore = new TestTaskStore();
+        var taskRegistryFactory = new LocalTaskRegistryFactory("user", taskStore);
+        var tracker = new TestProgressTracker(task, new GdsTestLog(), new Concurrency(4), taskRegistryFactory);
+
+        cnar.compute(tinyGraph, tracker);
+
+        assertThat(taskStore.tasks()).isEmpty();
+    }
 }

graph-sampling/src/test/java/org/neo4j/gds/graphsampling/samplers/rw/rwr/RandomWalkWithRestartsTest.java

@@ -21,16 +21,24 @@
 
 import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.neo4j.gds.NodeLabel;
+import org.neo4j.gds.TestProgressTracker;
+import org.neo4j.gds.TestTaskStore;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.GraphStore;
+import org.neo4j.gds.core.concurrency.Concurrency;
+import org.neo4j.gds.core.utils.progress.LocalTaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.extension.GdlExtension;
 import org.neo4j.gds.extension.GdlGraph;
 import org.neo4j.gds.extension.IdFunction;
 import org.neo4j.gds.extension.Inject;
 import org.neo4j.gds.graphsampling.config.RandomWalkWithRestartsConfig;
 import org.neo4j.gds.graphsampling.config.RandomWalkWithRestartsConfigImpl;
+import org.neo4j.gds.logging.GdsTestLog;
 import org.neo4j.gds.termination.TerminatedException;
 import org.neo4j.gds.termination.TerminationFlag;
 
@@ -79,6 +87,15 @@ class RandomWalkWithRestartsTest {
     @Inject
     private GraphStore graphStore;
 
+    @GdlGraph(graphNamePrefix = "tiny")
+    private static String TINY_GRAPH = "()-->()";
+
+    @Inject
+    private Graph tinyGraph;
+
+    @Inject
+    private GraphStore tinyGraphStore;
+
     Graph getGraph(RandomWalkWithRestartsConfig config) {
         return graphStore.getGraph(
             config.nodeLabelIdentifiers(graphStore),
@@ -400,4 +417,24 @@ void checkTerminationFlag() {
             .isInstanceOf(TerminatedException.class)
             .hasMessageContaining("The execution has been terminated.");
     }
+
+    @ParameterizedTest
+    @ValueSource(booleans = {true, false})
+    void progressTrackingTinyGraph(boolean nodeLabelStratification) {
+        var config = RandomWalkWithRestartsConfigImpl.builder()
+            .nodeLabelStratification(nodeLabelStratification)
+            .build();
+
+        var rwr = new RandomWalkWithRestarts(config);
+        Task task = rwr.progressTask(tinyGraphStore);
+
+        TestTaskStore taskStore = new TestTaskStore();
+        var taskRegistryFactory = new LocalTaskRegistryFactory("user", taskStore);
+        var tracker = new TestProgressTracker(task, new GdsTestLog(), new Concurrency(4), taskRegistryFactory);
+
+        rwr.compute(tinyGraph, tracker);
+
+        assertThat(taskStore.tasks()).isEmpty();
+    }
+
 }