Move up use of KnnContext

Author: jjaderberg

algo/src/main/java/org/neo4j/gds/similarity/filteredknn/FilteredKnn.java

@@ -38,6 +38,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.SplittableRandom;
+import java.util.concurrent.ExecutorService;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
@@ -50,7 +51,7 @@
 public class FilteredKnn extends Algorithm<FilteredKnn.Result> {
     private final Graph graph;
     private final FilteredNeighborFilterFactory neighborFilterFactory;
-    private final FilteredKnnContext context;
+    private final ExecutorService executorService;
     private final SplittableRandom splittableRandom;
     private final SimilarityComputer similarityComputer;
     private final List<Long> sourceNodes;
@@ -91,7 +92,7 @@ public static FilteredKnn create(
             sourceNodes,
             similarityComputer,
             neighborFilterFactory,
-            context,
+            context.executor(),
             splittableRandom,
             config.sampleRate(),
             config.deltaThreshold(),
@@ -134,7 +135,7 @@ private static SplittableRandom getSplittableRandom(Optional<Long> randomSeed) {
         List<Long> sourceNodes,
         SimilarityComputer similarityComputer,
         FilteredNeighborFilterFactory neighborFilterFactory,
-        FilteredKnnContext context,
+        ExecutorService executorService,
         SplittableRandom splittableRandom,
         double sampleRate,
         double deltaThreshold,
@@ -162,7 +163,7 @@ private static SplittableRandom getSplittableRandom(Optional<Long> randomSeed) {
         this.maxIterations = maxIterations;
         this.similarityComputer = similarityComputer;
         this.neighborFilterFactory = neighborFilterFactory;
-        this.context = context;
+        this.executorService = executorService;
         this.splittableRandom = splittableRandom;
         this.sourceNodes = sourceNodes;
         this.samplerSupplier = samplerSupplier;
@@ -172,8 +173,8 @@ public long nodeCount() {
         return graph.nodeCount();
     }
 
-    public FilteredKnnContext context() {
-        return context;
+    public ExecutorService executorService() {
+        return this.executorService;
     }
 
     @Override
@@ -218,7 +219,7 @@ public Result compute() {
                     ),
                     Optional.of(this.minBatchSize)
                 );
-                ParallelUtil.runWithConcurrency(this.concurrency, neighborFilterTasks, context.executor());
+                ParallelUtil.runWithConcurrency(this.concurrency, neighborFilterTasks, this.executorService);
             }
             this.progressTracker.endSubTask();
 
@@ -262,7 +263,7 @@ public void release() {
             Optional.of(this.minBatchSize)
         );
 
-        ParallelUtil.runWithConcurrency(this.concurrency, randomNeighborGenerators, context.executor());
+        ParallelUtil.runWithConcurrency(this.concurrency, randomNeighborGenerators, this.executorService);
 
         this.nodePairsConsidered += randomNeighborGenerators.stream().mapToLong(FilteredGenerateRandomNeighbors::neighborsFound).sum();
 
@@ -310,15 +311,12 @@ private long iteration(HugeObjectArray<FilteredNeighborList> neighbors) {
             return FilteredNeighborList.NOT_INSERTED;
         }
 
-        var executor = this.context.executor();
-
-
         // TODO: init in ctor and reuse - benchmark against new allocations
         var allOldNeighbors = HugeObjectArray.newArray(LongArrayList.class, nodeCount);
         var allNewNeighbors = HugeObjectArray.newArray(LongArrayList.class, nodeCount);
 
         progressTracker.beginSubTask();
-        ParallelUtil.readParallel(this.concurrency, nodeCount, executor, new FilteredSplitOldAndNewNeighbors(
+        ParallelUtil.readParallel(this.concurrency, nodeCount, this.executorService, new FilteredSplitOldAndNewNeighbors(
             this.splittableRandom,
             neighbors,
             allOldNeighbors,
@@ -369,7 +367,7 @@ private long iteration(HugeObjectArray<FilteredNeighborList> neighbors) {
         );
 
         progressTracker.beginSubTask();
-        ParallelUtil.runWithConcurrency(this.concurrency, neighborsJoiners, executor);
+        ParallelUtil.runWithConcurrency(this.concurrency, neighborsJoiners, executorService);
         progressTracker.endSubTask();
 
         this.nodePairsConsidered += neighborsJoiners.stream().mapToLong(JoinNeighbors::nodePairsConsidered).sum();

algo/src/main/java/org/neo4j/gds/similarity/knn/Knn.java

@@ -37,6 +37,7 @@
 
 import java.util.Optional;
 import java.util.SplittableRandom;
+import java.util.concurrent.ExecutorService;
 import java.util.function.Function;
 import java.util.function.UnaryOperator;
 import java.util.stream.IntStream;
@@ -49,7 +50,7 @@ public class Knn extends Algorithm<Knn.Result> {
     private final Graph graph;
     private final KnnBaseConfig config;
     private final NeighborFilterFactory neighborFilterFactory;
-    private final KnnContext context;
+    private final ExecutorService executorService;
     private final SplittableRandom splittableRandom;
     private final SimilarityComputer similarityComputer;
 
@@ -62,7 +63,7 @@ public static Knn createWithDefaults(Graph graph, KnnBaseConfig config, KnnConte
             config,
             SimilarityComputer.ofProperties(graph, config.nodeProperties()),
             new KnnNeighborFilterFactory(graph.nodeCount()),
-            context,
+            context.executor(),
             getSplittableRandom(config.randomSeed())
         );
     }
@@ -81,7 +82,7 @@ public static Knn create(
             config,
             similarityComputer,
             neighborFilterFactory,
-            context,
+            context.executor(),
             splittableRandom
         );
     }
@@ -97,24 +98,24 @@ private static SplittableRandom getSplittableRandom(Optional<Long> randomSeed) {
         KnnBaseConfig config,
         SimilarityComputer similarityComputer,
         NeighborFilterFactory neighborFilterFactory,
-        KnnContext context,
+        ExecutorService executorService,
         SplittableRandom splittableRandom
     ) {
         super(progressTracker);
         this.graph = graph;
         this.config = config;
         this.similarityComputer = similarityComputer;
         this.neighborFilterFactory = neighborFilterFactory;
-        this.context = context;
+        this.executorService = executorService;
         this.splittableRandom = splittableRandom;
     }
 
     public long nodeCount() {
         return graph.nodeCount();
     }
 
-    public KnnContext context() {
-        return context;
+    public ExecutorService executorService() {
+        return this.executorService;
     }
 
     @Override
@@ -161,7 +162,7 @@ public Result compute() {
                     ),
                     Optional.of(config.minBatchSize())
                 );
-                ParallelUtil.runWithConcurrency(config.concurrency(), neighborFilterTasks, context.executor());
+                ParallelUtil.runWithConcurrency(config.concurrency(), neighborFilterTasks, this.executorService);
             }
             this.progressTracker.endSubTask();
 
@@ -207,7 +208,7 @@ public void release() {
             Optional.of(config.minBatchSize())
         );
 
-        ParallelUtil.runWithConcurrency(config.concurrency(), randomNeighborGenerators, context.executor());
+        ParallelUtil.runWithConcurrency(config.concurrency(), randomNeighborGenerators, this.executorService);
 
         this.nodePairsConsidered += randomNeighborGenerators.stream().mapToLong(GenerateRandomNeighbors::neighborsFound).sum();
 
@@ -242,7 +243,6 @@ private long iteration(HugeObjectArray<NeighborList> neighbors) {
         }
 
         var concurrency = this.config.concurrency();
-        var executor = this.context.executor();
 
         var sampledK = this.config.sampledK(nodeCount);
 
@@ -251,7 +251,7 @@ private long iteration(HugeObjectArray<NeighborList> neighbors) {
         var allNewNeighbors = HugeObjectArray.newArray(LongArrayList.class, nodeCount);
 
         progressTracker.beginSubTask();
-        ParallelUtil.readParallel(concurrency, nodeCount, executor, new SplitOldAndNewNeighbors(
+        ParallelUtil.readParallel(concurrency, nodeCount, this.executorService, new SplitOldAndNewNeighbors(
             this.splittableRandom,
             neighbors,
             allOldNeighbors,
@@ -301,7 +301,7 @@ private long iteration(HugeObjectArray<NeighborList> neighbors) {
         );
 
         progressTracker.beginSubTask();
-        ParallelUtil.runWithConcurrency(concurrency, neighborsJoiners, executor);
+        ParallelUtil.runWithConcurrency(concurrency, neighborsJoiners, this.executorService);
         progressTracker.endSubTask();
 
         this.nodePairsConsidered += neighborsJoiners.stream().mapToLong(JoinNeighbors::nodePairsConsidered).sum();

proc/similarity/src/main/java/org/neo4j/gds/similarity/knn/KnnMutateProc.java

@@ -118,7 +118,7 @@ public ComputationResultConsumer<Knn, Knn.Result, KnnMutateConfig, Stream<Result
                     algorithm.nodeCount(),
                     config.concurrency(),
                     Objects.requireNonNull(result),
-                    algorithm.context()
+                    algorithm.executorService()
                 );
             }
 

proc/similarity/src/main/java/org/neo4j/gds/similarity/knn/KnnStatsProc.java

@@ -126,7 +126,7 @@ public Stream<Result> stats(ComputationResult<Knn, Knn.Result, KnnStatsConfig> c
                         algorithm.nodeCount(),
                         config.concurrency(),
                         result,
-                        algorithm.context()
+                        algorithm.executorService()
                     );
 
                     Graph similarityGraph = similarityGraphResult.similarityGraph();

proc/similarity/src/main/java/org/neo4j/gds/similarity/knn/KnnWriteProc.java

@@ -36,6 +36,7 @@
 
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ExecutorService;
 import java.util.stream.Stream;
 
 import static org.neo4j.gds.executor.ExecutionMode.WRITE_RELATIONSHIP;
@@ -100,7 +101,7 @@ protected SimilarityGraphResult similarityGraphResult(ComputationResult<Knn, Knn
             algorithm.nodeCount(),
             config.concurrency(),
             Objects.requireNonNull(computationResult.result()),
-            algorithm.context()
+            algorithm.executorService()
         );
     }
 
@@ -109,12 +110,12 @@ static SimilarityGraphResult computeToGraph(
         long nodeCount,
         int concurrency,
         Knn.Result result,
-        KnnContext context
+        ExecutorService executor
     ) {
         Graph similarityGraph = new SimilarityGraphBuilder(
             graph,
             concurrency,
-            context.executor()
+            executor
         ).build(result.streamSimilarityResult());
         return new SimilarityGraphResult(similarityGraph, nodeCount, false);
     }