Cluster hierarchy progress tracking

Co-authored-by: Veselin Nikolov <veselin.nikolov@neotechnology.com>

Author: IoannisPanagiotas

algo/src/main/java/org/neo4j/gds/hdbscan/ClusterHierarchy.java

@@ -22,6 +22,7 @@
 import org.neo4j.gds.collections.ha.HugeDoubleArray;
 import org.neo4j.gds.collections.ha.HugeLongArray;
 import org.neo4j.gds.collections.ha.HugeObjectArray;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 import java.util.function.Function;
 
@@ -49,7 +50,7 @@ final class ClusterHierarchy {
         this.nodeCount = nodeCount;
     }
 
-    static ClusterHierarchy create(long nodeCount, HugeObjectArray<Edge> edges) {
+    static ClusterHierarchy create(long nodeCount, HugeObjectArray<Edge> edges, ProgressTracker progressTracker) {
         var left = HugeLongArray.newArray(nodeCount);
         var right = HugeLongArray.newArray(nodeCount);
         var lambda = HugeDoubleArray.newArray(nodeCount);
@@ -61,6 +62,7 @@ static ClusterHierarchy create(long nodeCount, HugeObjectArray<Edge> edges) {
 
         var sizeFn = (Function<Long, Long>) n -> n < nodeCount ? 1L : size.get(n - nodeCount);
 
+        progressTracker.beginSubTask();
         for (var i = 0; i < edges.size(); i++) {
             var edge = edges.get(i);
             var l = unionFind.find(edge.source());
@@ -76,8 +78,10 @@ static ClusterHierarchy create(long nodeCount, HugeObjectArray<Edge> edges) {
             var rightSize = sizeFn.apply(r);
 
             size.set(adaptedIndex, leftSize + rightSize);
-        }
 
+            progressTracker.logProgress();
+        }
+        progressTracker.endSubTask();;
         return new ClusterHierarchy(currentRoot, left, right, lambda, size, nodeCount);
     }
 

algo/src/main/java/org/neo4j/gds/hdbscan/HDBScan.java

@@ -100,7 +100,7 @@ DualTreeMSTResult dualTreeMSTPhase(KdTree kdTree, CoreResult coreResult) {
     ClusterHierarchy createClusterHierarchy(DualTreeMSTResult dualTreeMSTResult){
         var edges = dualTreeMSTResult.edges();
         HugeSerialObjectMergeSort.sort(Edge.class, edges, Edge::distance);
-        return ClusterHierarchy.create(nodes.nodeCount(),edges);
+        return ClusterHierarchy.create(nodes.nodeCount(),edges,progressTracker);
     }
 
     KdTree buildKDTree() {

algo/src/main/java/org/neo4j/gds/hdbscan/HDBScanProgressTrackerCreator.java

@@ -25,7 +25,11 @@
 public class HDBScanProgressTrackerCreator {
 
     static Task kdBuildingTask(String name, long nodeCount){
-        return Tasks.leaf(name,nodeCount);
+        return Tasks.leaf(name, nodeCount);
+    }
+
+    static Task hierarchyTask(String name, long nodeCount){
+        return Tasks.leaf(name,nodeCount-1);
     }
 
 }

algo/src/test/java/org/neo4j/gds/hdbscan/ClusterHierarchyTest.java

@@ -19,12 +19,23 @@
  */
 package org.neo4j.gds.hdbscan;
 
+import org.assertj.core.api.Assertions;
 import org.assertj.core.api.SoftAssertions;
 import org.assertj.core.api.junit.jupiter.SoftAssertionsExtension;
 import org.assertj.core.data.Offset;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.neo4j.gds.collections.ha.HugeObjectArray;
+import org.neo4j.gds.compat.TestLog;
+import org.neo4j.gds.core.concurrency.Concurrency;
+import org.neo4j.gds.core.utils.logging.LoggerForProgressTrackingAdapter;
+import org.neo4j.gds.core.utils.progress.EmptyTaskRegistryFactory;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
+import org.neo4j.gds.core.utils.progress.tasks.TaskProgressTracker;
+import org.neo4j.gds.logging.GdsTestLog;
+
+import static org.neo4j.gds.assertj.Extractors.removingThreadId;
+import static org.neo4j.gds.assertj.Extractors.replaceTimings;
 
 @ExtendWith(SoftAssertionsExtension.class)
 class ClusterHierarchyTest {
@@ -36,7 +47,11 @@ void shouldWorkWithLineGraph(SoftAssertions assertions) {
             new Edge(0, 1, 5.)
         );
 
-        var clusterHierarchy = ClusterHierarchy.create(3, edges);
+        var clusterHierarchy = ClusterHierarchy.create(
+            3,
+            edges,
+            ProgressTracker.NULL_TRACKER
+        );
 
         // 1. `1` and `2` are joined and create new id = 3 --> first set
         // 2. `0` and `3` are joined and create new id = 4 --> second set
@@ -71,7 +86,11 @@ void shouldWorkOnMoreComplexTree(SoftAssertions assertions) {
             new Edge(5, 7, 1.42823558)
         );
 
-        var clusterHierarchy = ClusterHierarchy.create(edges.size() + 1, edges);
+        var clusterHierarchy = ClusterHierarchy.create(
+            edges.size() + 1,
+            edges,
+            ProgressTracker.NULL_TRACKER
+        );
 
 
         assertions.assertThat(clusterHierarchy.left(11)).isEqualTo(2);
@@ -124,4 +143,31 @@ void shouldWorkOnMoreComplexTree(SoftAssertions assertions) {
         assertions.assertThat(clusterHierarchy.lambda(20)).isCloseTo(1.42823558, Offset.offset(1e-9));
         assertions.assertThat(clusterHierarchy.size(20)).isEqualTo(11);
     }
+
+    @Test
+    void shouldLogProgress(){
+        var edges = HugeObjectArray.of(
+            new Edge(1, 2, 3.),
+            new Edge(0, 1, 5.)
+        );
+
+        var progressTask = HDBScanProgressTrackerCreator.hierarchyTask("foo",3);
+        var log = new GdsTestLog();
+        var progressTracker = new TaskProgressTracker(progressTask, new LoggerForProgressTrackingAdapter(log), new Concurrency(1), EmptyTaskRegistryFactory.INSTANCE);
+        var clusterHierarchy = ClusterHierarchy.create(
+            3,
+            edges,
+            progressTracker
+        );
+
+        Assertions.assertThat(log.getMessages(TestLog.INFO))
+            .extracting(removingThreadId())
+            .extracting(replaceTimings())
+            .containsExactly(
+                "foo :: Start",
+                "foo 50%",
+                "foo 100%",
+                "foo :: Finished"
+            );
+    }
 }