KnnContext is a record

Author: jjaderberg

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

@@ -19,27 +19,21 @@
  */
 package org.neo4j.gds.similarity.knn;
 
-import org.immutables.value.Value;
-import org.neo4j.gds.annotation.ValueClass;
 import org.neo4j.gds.core.concurrency.DefaultPool;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 
 import java.util.concurrent.ExecutorService;
 
-@ValueClass
-public interface KnnContext {
+public record KnnContext(ExecutorService executor, ProgressTracker progressTracker) {
 
-    @Value.Default
-    default ExecutorService executor() {
-        return DefaultPool.INSTANCE;
+    public KnnContext(ProgressTracker progressTracker) {
+        this(DefaultPool.INSTANCE, progressTracker);
     }
 
-    @Value.Default
-    default ProgressTracker progressTracker() {
-        return ProgressTracker.NULL_TRACKER;
-    }
-
-    static KnnContext empty() {
-        return ImmutableKnnContext.builder().build();
+    public static KnnContext empty() {
+        return new KnnContext(
+            DefaultPool.INSTANCE,
+            ProgressTracker.NULL_TRACKER
+        );
     }
 }

algo/src/test/java/org/neo4j/gds/similarity/filteredknn/FilteredKnnTest.java

@@ -31,7 +31,6 @@
 import org.neo4j.gds.similarity.NodeFilterSpec;
 import org.neo4j.gds.similarity.SimilarityResult;
 import org.neo4j.gds.similarity.filtering.NodeIdNodeFilterSpec;
-import org.neo4j.gds.similarity.knn.ImmutableKnnContext;
 import org.neo4j.gds.similarity.knn.KnnContext;
 import org.neo4j.gds.similarity.knn.KnnNodePropertySpec;
 import org.neo4j.gds.similarity.knn.KnnParametersSansNodeCount;
@@ -108,7 +107,7 @@ void shouldRunJustLikeKnnWhenYouDoNotSpecifySourceNodeFilterOrTargetNodeFilter()
         var params = filteredSans.finalize(graph.nodeCount());
 
 
-        var knnContext = ImmutableKnnContext.builder().build();
+        var knnContext = KnnContext.empty();
 
         var knn = FilteredKnn.createWithoutSeeding(graph, params, knnContext, TerminationFlag.RUNNING_TRUE);
         var result = knn.compute();

algo/src/test/java/org/neo4j/gds/similarity/knn/KnnTest.java

@@ -34,7 +34,6 @@
 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.concurrency.DefaultPool;
 import org.neo4j.gds.core.loading.NullPropertyMap;
 import org.neo4j.gds.extension.GdlExtension;
 import org.neo4j.gds.extension.GdlGraph;
@@ -95,7 +94,7 @@ class KnnTest {
     @Inject
     private TestGraph multPropMissingGraph;
 
-    private KnnContext context = ImmutableKnnContext.builder().executor(DefaultPool.INSTANCE).build();
+    private final KnnContext context = KnnContext.empty();
 
     @Test
     void shouldRun() {
@@ -559,10 +558,7 @@ void shouldLogProgress() {
             new Concurrency(1)
         );
 
-        var progressContext = ImmutableKnnContext.builder()
-            .progressTracker(progressTrackerWithLog.progressTracker())
-            .executor(DefaultPool.INSTANCE)
-            .build();
+        var progressContext = new KnnContext(progressTrackerWithLog.progressTracker());
         var log = progressTrackerWithLog.log();
 
         var knn = Knn.create(

applications/algorithms/similarity/src/main/java/org/neo4j/gds/applications/algorithms/similarity/SimilarityAlgorithms.java

@@ -27,8 +27,8 @@
 import org.neo4j.gds.similarity.filteredknn.FilteredKnnResult;
 import org.neo4j.gds.similarity.filterednodesim.FilteredNodeSimilarityParameters;
 import org.neo4j.gds.similarity.filtering.NodeFilter;
-import org.neo4j.gds.similarity.knn.ImmutableKnnContext;
 import org.neo4j.gds.similarity.knn.Knn;
+import org.neo4j.gds.similarity.knn.KnnContext;
 import org.neo4j.gds.similarity.knn.KnnNeighborFilterFactory;
 import org.neo4j.gds.similarity.knn.KnnParameters;
 import org.neo4j.gds.similarity.knn.KnnResult;
@@ -55,11 +55,7 @@ public FilteredKnnResult filteredKnn(
         FilteredKnnParameters parameters,
         ProgressTracker progressTracker
     ) {
-        var knnContext = ImmutableKnnContext
-            .builder()
-            .progressTracker(progressTracker)
-            .executor(DefaultPool.INSTANCE)
-            .build();
+        var knnContext = new KnnContext(progressTracker);
 
         var algorithm = FilteredKNNFactory.create(graph, parameters, knnContext,terminationFlag);
 
@@ -98,11 +94,7 @@ KnnResult knn(Graph graph, KnnParameters parameters, ProgressTracker progressTra
             new SimilarityFunction(SimilarityComputer.ofProperties(graph, parameters.nodePropertySpecs())),
             new KnnNeighborFilterFactory(graph.nodeCount()),
             Optional.empty(),
-            ImmutableKnnContext
-                .builder()
-                .progressTracker(progressTracker)
-                .executor(DefaultPool.INSTANCE)
-                .build(),
+            new KnnContext(progressTracker),
             terminationFlag
         );
 

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/ApproximateLinkPrediction.java

@@ -20,15 +20,14 @@
 package org.neo4j.gds.procedures.pipelines;
 
 import org.neo4j.gds.api.Graph;
-import org.neo4j.gds.core.concurrency.DefaultPool;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.mem.MemoryEstimation;
 import org.neo4j.gds.ml.linkmodels.LinkPredictionResult;
 import org.neo4j.gds.ml.linkmodels.PredictedLink;
 import org.neo4j.gds.ml.models.Classifier;
 import org.neo4j.gds.ml.pipeline.linkPipeline.LinkFeatureExtractor;
-import org.neo4j.gds.similarity.knn.ImmutableKnnContext;
 import org.neo4j.gds.similarity.knn.Knn;
+import org.neo4j.gds.similarity.knn.KnnContext;
 import org.neo4j.gds.similarity.knn.KnnParameters;
 import org.neo4j.gds.similarity.knn.KnnResult;
 import org.neo4j.gds.termination.TerminationFlag;
@@ -80,8 +79,7 @@ LinkPredictionResult predictLinks(LinkPredictionSimilarityComputer linkPredictio
                 targetNodeFilter
             ),
             Optional.empty(),
-            ImmutableKnnContext.of(
-                DefaultPool.INSTANCE,
+            new KnnContext(
                 progressTracker
             ),
             TerminationFlag.RUNNING_TRUE