Node2Vec also gets walkBufferSize

Author: jjaderberg

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/Node2Vec.java

@@ -47,6 +47,7 @@ public class Node2Vec extends Algorithm<Node2VecModel.Result> {
     private final List<Long> sourceNodes;
     private final Optional<Long> maybeRandomSeed;
     private final TrainParameters trainParameters;
+    private final int walkBufferSize;
 
 
     public static MemoryEstimation memoryEstimation(int walksPerNode, int walkLength, int embeddingDimension) {
@@ -84,6 +85,7 @@ static Node2Vec create(
             concurrency,
             List.of(),
             maybeRandomSeed,
+            1000,
             walkParameters,
             trainParameters,
             progressTracker
@@ -95,6 +97,7 @@ public Node2Vec(
         int concurrency,
         List<Long> sourceNodes,
         Optional<Long> maybeRandomSeed,
+        int walkBufferSize,
         WalkParameters walkParameters,
         TrainParameters trainParameters,
         ProgressTracker progressTracker
@@ -103,6 +106,7 @@ public Node2Vec(
         this.graph = graph;
         this.concurrency = concurrency;
         this.walkParameters = walkParameters;
+        this.walkBufferSize = walkBufferSize;
         this.sourceNodes = sourceNodes;
         this.maybeRandomSeed = maybeRandomSeed;
         this.trainParameters = trainParameters;
@@ -140,6 +144,7 @@ public Node2VecModel.Result compute() {
             concurrency,
             sourceNodes,
             walkParameters,
+            walkBufferSize,
             DefaultPool.INSTANCE,
             progressTracker,
             terminationFlag
@@ -185,6 +190,7 @@ private List<Node2VecRandomWalkTask> walkTasks(
         int concurrency,
         List<Long> sourceNodes,
         WalkParameters walkParameters,
+        int walkBufferSize,
         ExecutorService executorService,
         ProgressTracker progressTracker,
         TerminationFlag terminationFlag
@@ -211,6 +217,7 @@ private List<Node2VecRandomWalkTask> walkTasks(
                 index,
                 compressedRandomWalks,
                 randomWalkPropabilitiesBuilder,
+                walkBufferSize,
                 randomSeed,
                 walkParameters.walkLength,
                 walkParameters.returnFactor,

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/Node2VecAlgorithmFactory.java

@@ -52,6 +52,7 @@ public Node2Vec build(
             configuration.concurrency(),
             configuration.sourceNodes(),
             configuration.randomSeed(),
+            configuration.walkBufferSize(),
             configuration.walkParameters(),
             configuration.trainParameters(),
             progressTracker

algo/src/main/java/org/neo4j/gds/embeddings/node2vec/Node2VecRandomWalkTask.java

@@ -38,6 +38,7 @@ final class Node2VecRandomWalkTask implements Runnable {
     private final CompressedRandomWalks compressedRandomWalks;
     private final RandomWalkProbabilities.Builder randomWalkProbabilitiesBuilder;
     private final RandomWalkSampler sampler;
+    private final int walkBufferSize;
     private int walks;
     private int maxWalkLength;
     private long maxIndex;
@@ -52,6 +53,7 @@ final class Node2VecRandomWalkTask implements Runnable {
         AtomicLong walkIndex,
         CompressedRandomWalks compressedRandomWalks,
         RandomWalkProbabilities.Builder randomWalkProbabilitiesBuilder,
+        int walkBufferSize,
         long randomSeed,
         int walkLength,
         double returnFactor,
@@ -65,6 +67,7 @@ final class Node2VecRandomWalkTask implements Runnable {
         this.walkIndex = walkIndex;
         this.compressedRandomWalks = compressedRandomWalks;
         this.randomWalkProbabilitiesBuilder = randomWalkProbabilitiesBuilder;
+        this.walkBufferSize = walkBufferSize;
 
         this.sampler = RandomWalkSampler.create(
             graph,
@@ -85,7 +88,7 @@ private boolean consumePath(long[] path) {
         randomWalkProbabilitiesBuilder.registerWalk(path);
         compressedRandomWalks.add(index, path);
         maxWalkLength = Math.max(path.length, maxWalkLength);
-        if (walks++ == 1000) { //this is just to get the same
+        if (walks++ == walkBufferSize) {
             walks = 0;
             return this.terminationFlag.running();
         }