Merge pull request #5211 from s1ck/pregel-expose-log-methods

Expose log methods in PregelContext

Author: s1ck

doc/asciidoc/algorithms/algorithms-pregel-api.adoc

@@ -349,6 +349,34 @@ public class CustomComputation implements PregelComputation<CustomConfig> {
 }
 ----
 
+[[algorithms-pregel-api-logging]]
+=== Logging
+
+The following methods are available for all contexts (`InitContext`, `ComputeContext`, `MasterComputeContext`) to inject custom messages into the progress log of the algorithm execution.
+
+.The log methods can be used in Pregel contexts
+[source, java]
+----
+// All contexts inherit from PregelContext
+public abstract class PregelContext<CONFIG extends PregelConfig> {
+
+    // Log a debug message to the Neo4j log.
+    public void logDebug(String message) {
+        progressTracker.logDebug(message);
+    }
+
+    // Log a warning message to the Neo4j log.
+    public void logWarning(String message) {
+        progressTracker.logWarning(message);
+    }
+
+    // Log a info message to the Neo4j log
+    public void logMessage(String message) {
+        progressTracker.logMessage(message);
+    }
+
+}
+----
 
 [[algorithms-pregel-api-procedure]]
 == Run Pregel via Cypher

pregel/src/main/java/org/neo4j/gds/beta/pregel/ForkJoinComputeStep.java

@@ -75,10 +75,10 @@ public final class ForkJoinComputeStep<CONFIG extends PregelConfig, ITERATOR ext
         this.nodeBatch = nodeBatch;
         this.nodeValue = nodeValue;
         this.messenger = messenger;
-        this.computeContext = new ComputeContext<>(this, config);
+        this.computeContext = new ComputeContext<>(this, config, progressTracker);
         this.sentMessage = sentMessage;
         this.progressTracker = progressTracker;
-        this.initContext = new InitContext<>(this, config, graph);
+        this.initContext = new InitContext<>(this, config, graph, progressTracker);
     }
 
     @Override

pregel/src/main/java/org/neo4j/gds/beta/pregel/PartitionedComputeStep.java

@@ -60,9 +60,9 @@ public final class PartitionedComputeStep<CONFIG extends PregelConfig, ITERATOR
         this.voteBits = voteBits;
         this.nodeBatch = nodeBatch;
         this.messenger = messenger;
-        this.computeContext = new ComputeContext<>(this, config);
+        this.computeContext = new ComputeContext<>(this, config, progressTracker);
         this.progressTracker = progressTracker;
-        this.initContext = new InitContext<>(this, config, graph);
+        this.initContext = new InitContext<>(this, config, graph, progressTracker);
     }
 
     @Override

pregel/src/main/java/org/neo4j/gds/beta/pregel/Pregel.java

@@ -200,7 +200,7 @@ public void release() {
     }
 
     private boolean runMasterComputeStep(int iteration) {
-        var context = new MasterComputeContext<>(config, graph, iteration, nodeValues, executor);
+        var context = new MasterComputeContext<>(config, graph, iteration, nodeValues, executor, progressTracker);
         var didConverge = computation.masterCompute(context);
         return didConverge;
     }

pregel/src/main/java/org/neo4j/gds/beta/pregel/context/ComputeContext.java

@@ -21,6 +21,7 @@
 
 import org.neo4j.gds.beta.pregel.ComputeStep;
 import org.neo4j.gds.beta.pregel.PregelConfig;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 /**
  * A context that is used during the computation. It allows an implementation
@@ -29,8 +30,8 @@
  */
 public final class ComputeContext<CONFIG extends PregelConfig> extends NodeCentricContext<CONFIG> {
 
-    public ComputeContext(ComputeStep<CONFIG, ?> computeStep, CONFIG config) {
-        super(computeStep, config);
+    public ComputeContext(ComputeStep<CONFIG, ?> computeStep, CONFIG config, ProgressTracker progressTracker) {
+        super(computeStep, config, progressTracker);
         this.sendMessagesFunction = config.hasRelationshipWeightProperty()
             ? computeStep::sendToNeighborsWeighted
             : computeStep::sendToNeighbors;

pregel/src/main/java/org/neo4j/gds/beta/pregel/context/InitContext.java

@@ -23,6 +23,7 @@
 import org.neo4j.gds.api.properties.nodes.NodePropertyValues;
 import org.neo4j.gds.beta.pregel.ComputeStep;
 import org.neo4j.gds.beta.pregel.PregelConfig;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 import java.util.Set;
 
@@ -35,8 +36,13 @@
 public final class InitContext<CONFIG extends PregelConfig> extends NodeCentricContext<CONFIG> {
     private final NodePropertyContainer nodePropertyContainer;
 
-    public InitContext(ComputeStep<CONFIG, ?> computeStep, CONFIG config, NodePropertyContainer nodePropertyContainer) {
-        super(computeStep, config);
+    public InitContext(
+        ComputeStep<CONFIG, ?> computeStep,
+        CONFIG config,
+        NodePropertyContainer nodePropertyContainer,
+        ProgressTracker progressTracker
+    ) {
+        super(computeStep, config, progressTracker);
         this.nodePropertyContainer = nodePropertyContainer;
     }
 

pregel/src/main/java/org/neo4j/gds/beta/pregel/context/MasterComputeContext.java

@@ -22,6 +22,7 @@
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.beta.pregel.NodeValue;
 import org.neo4j.gds.beta.pregel.PregelConfig;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 import java.util.concurrent.ExecutorService;
 import java.util.function.LongPredicate;
@@ -39,9 +40,10 @@ public MasterComputeContext(
         Graph graph,
         int iteration,
         NodeValue nodeValue,
-        ExecutorService executorService
+        ExecutorService executorService,
+        ProgressTracker progressTracker
     ) {
-        super(config);
+        super(config, progressTracker);
         this.graph = graph;
         this.iteration = iteration;
         this.nodeValue = nodeValue;

pregel/src/main/java/org/neo4j/gds/beta/pregel/context/NodeCentricContext.java

@@ -21,6 +21,7 @@
 
 import org.neo4j.gds.beta.pregel.ComputeStep;
 import org.neo4j.gds.beta.pregel.PregelConfig;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 import java.util.function.LongConsumer;
 
@@ -30,8 +31,8 @@ public abstract class NodeCentricContext<CONFIG extends PregelConfig> extends Pr
 
     long nodeId;
 
-    NodeCentricContext(ComputeStep<CONFIG, ?> computeStep, CONFIG config) {
-        super(config);
+    NodeCentricContext(ComputeStep<CONFIG, ?> computeStep, CONFIG config, ProgressTracker progressTracker) {
+        super(config, progressTracker);
         this.computeStep = computeStep;
     }
 

pregel/src/main/java/org/neo4j/gds/beta/pregel/context/PregelContext.java

@@ -20,13 +20,16 @@
 package org.neo4j.gds.beta.pregel.context;
 
 import org.neo4j.gds.beta.pregel.PregelConfig;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 public abstract class PregelContext<CONFIG extends PregelConfig> {
 
     protected final CONFIG config;
+    private final ProgressTracker progressTracker;
 
-    protected PregelContext(CONFIG config) {
+    protected PregelContext(CONFIG config, ProgressTracker progressTracker) {
         this.config = config;
+        this.progressTracker = progressTracker;
     }
 
     /**
@@ -36,6 +39,27 @@ public CONFIG config() {
         return config;
     }
 
+    /**
+     * Log a debug message to the Neo4j log.
+     */
+    public void logDebug(String message) {
+        progressTracker.logDebug(message);
+    }
+
+    /**
+     * Log a warning message to the Neo4j log.
+     */
+    public void logWarning(String message) {
+        progressTracker.logWarning(message);
+    }
+
+    /**
+     * Log a info message to the Neo4j log
+     */
+    public void logMessage(String message) {
+        progressTracker.logMessage(message);
+    }
+
     /**
      * Indicates whether the input graph is a multi-graph.
      */