Estimate relationship counts at creation of some progress tasks in LP

Author: breakanalysis

pipeline/src/main/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrain.java

@@ -95,19 +95,31 @@ public LinkPredictionTrain(
         this.classIdMap = makeClassIdMap();
     }
 
-    public static List<Task> progressTasks(int validationFolds, int numberOfModelSelectionTrials) {
+    public static List<Task> progressTasks(
+        long relationshipCount,
+        LinkPredictionSplitConfig splitConfig,
+        int numberOfModelSelectionTrials
+    ) {
+        // the relationship count estimates depend on both UndirectedEdgeSplitter
+        // and the volume set in extractFeaturesAndLabels
+        var selectionRatio = (1 + splitConfig.negativeSamplingRatio());
+        double nonTestRelationshipCount = relationshipCount * (1 - splitConfig.testFraction());
         return List.of(
-            Tasks.leaf("Extract train features"),
+            Tasks.leaf("Extract train features",
+                (long) (nonTestRelationshipCount * splitConfig.trainFraction() * selectionRatio)
+            ),
             Tasks.iterativeFixed(
                 "Select best model",
-                () -> List.of(Tasks.leaf("Trial", validationFolds)),
+                () -> List.of(Tasks.leaf("Trial", splitConfig.validationFolds())),
                 numberOfModelSelectionTrials
             ),
             ClassifierTrainer.progressTask("Train best model"),
             Tasks.leaf("Compute train metrics"),
             Tasks.task(
                 "Evaluate on test data",
-                Tasks.leaf("Extract test features"),
+                Tasks.leaf("Extract test features",
+                    (long) (relationshipCount * splitConfig.testFraction() * selectionRatio)
+                ),
                 Tasks.leaf("Compute test metrics")
             )
         );

pipeline/src/test/java/org/neo4j/gds/ml/pipeline/linkPipeline/train/LinkPredictionTrainTest.java

@@ -461,7 +461,8 @@ void logProgressRF() {
         var log = Neo4jProxy.testLog();
         var progressTracker = new TestProgressTracker(
             progressTask(
-                pipeline.splitConfig().validationFolds(),
+                trainGraph.relationshipCount(),
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ),
             log,
@@ -552,7 +553,8 @@ void logProgressLR() {
         var log = Neo4jProxy.testLog();
         var progressTracker = new TestProgressTracker(
             progressTask(
-                pipeline.splitConfig().validationFolds(),
+                trainGraph.relationshipCount(),
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ),
             log,
@@ -614,7 +616,8 @@ void logProgressLRWithRange() {
         var log = Neo4jProxy.testLog();
         var progressTracker = new TestProgressTracker(
             progressTask(
-                pipeline.splitConfig().validationFolds(),
+                trainGraph.relationshipCount(),
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ),
             log,
@@ -713,8 +716,11 @@ void logProgressLRWithRange() {
             );
     }
 
-    static Task progressTask(int validationFolds, int numberOfModelSelectionTrials) {
-        return Tasks.task("MY TEST TASK", LinkPredictionTrain.progressTasks(validationFolds, numberOfModelSelectionTrials));
+    static Task progressTask(long relationshipCount, LinkPredictionSplitConfig splitConfig, int numberOfModelSelectionTrials) {
+        return Tasks.task(
+            "MY TEST TASK",
+            LinkPredictionTrain.progressTasks(relationshipCount, splitConfig, numberOfModelSelectionTrials)
+        );
     }
 
     private LinkPredictionTrainingPipeline linkPredictionPipeline() {

proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineAlgorithmFactory.java

@@ -70,9 +70,15 @@ public String taskName() {
 
     @Override
     public Task progressTask(GraphStore graphStore, LinkPredictionTrainConfig config) {
+        var relationshipCount = config
+            .internalRelationshipTypes(graphStore)
+            .stream()
+            .mapToLong(graphStore::relationshipCount)
+            .sum();
         return LinkPredictionTrainPipelineExecutor.progressTask(
             taskName(),
-            PipelineCatalog.getTyped(config.username(), config.pipeline(), LinkPredictionTrainingPipeline.class)
+            PipelineCatalog.getTyped(config.username(), config.pipeline(), LinkPredictionTrainingPipeline.class),
+            relationshipCount
         );
     }
 

proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineExecutor.java

@@ -81,7 +81,7 @@ public LinkPredictionTrainPipelineExecutor(
         );
     }
 
-    public static Task progressTask(String taskName, LinkPredictionTrainingPipeline pipeline) {
+    public static Task progressTask(String taskName, LinkPredictionTrainingPipeline pipeline, long relationshipCount) {
         return Tasks.task(taskName, new ArrayList<>() {{
             add(Tasks.leaf("Split relationships"));
             add(Tasks.iterativeFixed(
@@ -90,7 +90,8 @@ public static Task progressTask(String taskName, LinkPredictionTrainingPipeline
                 pipeline.nodePropertySteps().size()
             ));
             addAll(LinkPredictionTrain.progressTasks(
-                pipeline.splitConfig().validationFolds(),
+                relationshipCount,
+                pipeline.splitConfig(),
                 pipeline.numberOfModelSelectionTrials()
             ));
         }});

proc/machine-learning/src/test/java/org/neo4j/gds/ml/linkmodels/pipeline/train/LinkPredictionTrainPipelineExecutorTest.java

@@ -386,8 +386,17 @@ void shouldLogProgress() {
 
         TestProcedureRunner.applyOnProcedure(db, TestProc.class, caller -> {
             var log = Neo4jProxy.testLog();
+            var relationshipCount = config
+                .internalRelationshipTypes(graphStore)
+                .stream()
+                .mapToLong(graphStore::relationshipCount)
+                .sum();
             var progressTracker = new TestProgressTracker(
-                LinkPredictionTrainPipelineExecutor.progressTask("Link Prediction Train Pipeline", pipeline),
+                LinkPredictionTrainPipelineExecutor.progressTask(
+                    "Link Prediction Train Pipeline",
+                    pipeline,
+                    relationshipCount
+                ),
                 log,
                 1,
                 EmptyTaskRegistryFactory.INSTANCE