Merge pull request #9550 from neo-technology/knn-task-label

Use typed task label for filtered KNN
neo4j · Aug 28, 2024 · 1ee4d03 · 1ee4d03
2 parents 12a20d6 + d1bb82e
commit 1ee4d03
Showing 1 changed file with 17 additions and 3 deletions.
diff --git a/.../src/main/java/org/neo4j/gds/applications/algorithms/similarity/SimilarityAlgorithms.java b/.../src/main/java/org/neo4j/gds/applications/algorithms/similarity/SimilarityAlgorithms.java
@@ -36,7 +36,6 @@
 import org.neo4j.gds.similarity.knn.Knn;
 import org.neo4j.gds.similarity.knn.KnnBaseConfig;
 import org.neo4j.gds.similarity.knn.KnnContext;
-import org.neo4j.gds.similarity.knn.KnnFactory;
 import org.neo4j.gds.similarity.knn.KnnNeighborFilterFactory;
 import org.neo4j.gds.similarity.knn.KnnResult;
 import org.neo4j.gds.similarity.knn.metrics.SimilarityComputer;
@@ -47,6 +46,7 @@
 
 import java.util.List;
 
+import static org.neo4j.gds.applications.algorithms.metadata.LabelForProgressTracking.FilteredKNN;
 import static org.neo4j.gds.applications.algorithms.metadata.LabelForProgressTracking.FilteredNodeSimilarity;
 import static org.neo4j.gds.applications.algorithms.metadata.LabelForProgressTracking.KNN;
 
@@ -65,10 +65,24 @@ public SimilarityAlgorithms(
     }
 
     FilteredKnnResult filteredKnn(Graph graph, FilteredKnnBaseConfig configuration) {
-        var taskTree = KnnFactory.knnTaskTree(graph.nodeCount(), configuration.maxIterations());
+        long nodeCount = graph.nodeCount();
+
+        Task task = Tasks.task(
+            FilteredKNN.value,
+            Tasks.leaf("Initialize random neighbors", nodeCount),
+            Tasks.iterativeDynamic(
+                "Iteration",
+                () -> List.of(
+                    Tasks.leaf("Split old and new neighbors", nodeCount),
+                    Tasks.leaf("Reverse old and new neighbors", nodeCount),
+                    Tasks.leaf("Join neighbors", nodeCount)
+                ),
+                configuration.maxIterations()
+            )
+        );
         var progressTracker = progressTrackerCreator.createProgressTracker(
             configuration,
-            taskTree
+            task
         );
         var knnContext = ImmutableKnnContext
             .builder()