add instruumentation to KNN so we can spy on visited neighbours

Author: lassewesth

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

@@ -39,6 +39,8 @@ final class GenerateRandomNeighbors implements Runnable {
     private final int boundedK;
     private final ProgressTracker progressTracker;
     private final Partition partition;
+    private final NeighbourConsumers neighbourConsumers;
+
     private long neighborsFound;
 
     GenerateRandomNeighbors(
@@ -50,7 +52,8 @@ final class GenerateRandomNeighbors implements Runnable {
         int k,
         int boundedK,
         Partition partition,
-        ProgressTracker progressTracker
+        ProgressTracker progressTracker,
+        NeighbourConsumers neighbourConsumers
     ) {
         this.sampler = sampler;
         this.random = random;
@@ -62,6 +65,7 @@ final class GenerateRandomNeighbors implements Runnable {
         this.progressTracker = progressTracker;
         this.partition = partition;
         this.neighborsFound = 0;
+        this.neighbourConsumers = neighbourConsumers;
     }
 
     @Override
@@ -80,7 +84,7 @@ public void run() {
                 l -> neighborFilter.excludeNodePair(nodeId, l)
             );
 
-            var neighbors = new NeighborList(k);
+            var neighbors = new NeighborList(k, neighbourConsumers.get(nodeId));
             for (long candidate : chosen) {
                 neighbors.add(candidate, computer.safeSimilarity(nodeId, candidate), rng, 0.0);
             }

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

@@ -53,18 +53,30 @@ public class Knn extends Algorithm<Knn.Result> {
     private final ExecutorService executorService;
     private final SplittableRandom splittableRandom;
     private final SimilarityComputer similarityComputer;
+    private final NeighbourConsumers neighborConsumers;
 
     private long nodePairsConsidered;
 
     public static Knn createWithDefaults(Graph graph, KnnBaseConfig config, KnnContext context) {
+        return createWithDefaultsss(graph, config, context, NeighbourConsumers.devNull);
+    }
+
+    @NotNull
+    public static Knn createWithDefaultsss(
+        Graph graph,
+        KnnBaseConfig config,
+        KnnContext context,
+        NeighbourConsumers neighborConsumers
+    ) {
         return new Knn(
             context.progressTracker(),
             graph,
             config,
             SimilarityComputer.ofProperties(graph, config.nodeProperties()),
             new KnnNeighborFilterFactory(graph.nodeCount()),
             context.executor(),
-            getSplittableRandom(config.randomSeed())
+            getSplittableRandom(config.randomSeed()),
+            neighborConsumers
         );
     }
 
@@ -83,7 +95,8 @@ public static Knn create(
             similarityComputer,
             neighborFilterFactory,
             context.executor(),
-            splittableRandom
+            splittableRandom,
+            NeighbourConsumers.devNull
         );
     }
 
@@ -99,7 +112,8 @@ private static SplittableRandom getSplittableRandom(Optional<Long> randomSeed) {
         SimilarityComputer similarityComputer,
         NeighborFilterFactory neighborFilterFactory,
         ExecutorService executorService,
-        SplittableRandom splittableRandom
+        SplittableRandom splittableRandom,
+        NeighbourConsumers neighborConsumers
     ) {
         super(progressTracker);
         this.graph = graph;
@@ -108,6 +122,7 @@ private static SplittableRandom getSplittableRandom(Optional<Long> randomSeed) {
         this.neighborFilterFactory = neighborFilterFactory;
         this.executorService = executorService;
         this.splittableRandom = splittableRandom;
+        this.neighborConsumers = neighborConsumers;
     }
 
     public long nodeCount() {
@@ -202,7 +217,8 @@ public void release() {
                     k,
                     boundedK,
                     partition,
-                    progressTracker
+                    progressTracker,
+                    neighborConsumers
                 );
             },
             Optional.of(config.minBatchSize())

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

@@ -81,6 +81,9 @@ static boolean isChecked(long value) {
     static final int NOT_INSERTED = 0;
     private static final int INSERTED = 1;
 
+    // a listener that sees every neighbour someone tries to add
+    private final NeighbourConsumer neighbourConsumer;
+
     // maximum number of elements, aka the top K
     private final int elementCapacity;
     // currently stored number of elements
@@ -90,13 +93,14 @@ static boolean isChecked(long value) {
     // every item occupies two entries in the array, [ doubleToLongBits(priority), element ]
     private final long[] priorityElementPairs;
 
-    NeighborList(int elementCapacity) {
+    NeighborList(int elementCapacity, NeighbourConsumer neighbourConsumer) {
         if (elementCapacity <= 0) {
             throw new IllegalArgumentException("Bound cannot be smaller than or equal to 0");
         }
 
         this.elementCapacity = elementCapacity;
         this.priorityElementPairs = new long[elementCapacity * 2];
+        this.neighbourConsumer = neighbourConsumer;
     }
 
     public LongStream elements() {
@@ -125,6 +129,8 @@ long getAndFlagAsChecked(int index) {
      * This allows KNN to just add the return values together without having the check on each of them.
      */
     public long add(long element, double priority, SplittableRandom random, double perturbationRate) {
+        neighbourConsumer.offer(element, priority);
+
         int insertIdx = 0;
         int currNumElementsWithPriority = elementCount * 2;
 

…edknn/FilteredNeighborFilterFactory.java …ds/similarity/knn/NeighbourConsumer.javaalgo/src/main/java/org/neo4j/gds/similarity/filteredknn/FilteredNeighborFilterFactory.java renamed to algo/src/main/java/org/neo4j/gds/similarity/knn/NeighbourConsumer.java algo/src/main/java/org/neo4j/gds/similarity/filteredknn/FilteredNeighborFilterFactory.java renamed to algo/src/main/java/org/neo4j/gds/similarity/knn/NeighbourConsumer.java

@@ -17,8 +17,13 @@
  * You should have received a copy of the GNU General Public License
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
-package org.neo4j.gds.similarity.filteredknn;
+package org.neo4j.gds.similarity.knn;
 
-public interface FilteredNeighborFilterFactory {
-    FilteredNeighborFilter create();
+/**
+ * During KNN execution, this listener is offered every neighbour we encounter.
+ */
+public interface NeighbourConsumer {
+    NeighbourConsumer devNull = (element, priority) -> { /* do nothing */ };
+
+    void offer(long element, double priority);
 }

…nn/FilteredKnnNeighborFilterFactory.java …s/similarity/knn/NeighbourConsumers.javaalgo/src/main/java/org/neo4j/gds/similarity/filteredknn/FilteredKnnNeighborFilterFactory.java renamed to algo/src/main/java/org/neo4j/gds/similarity/knn/NeighbourConsumers.java algo/src/main/java/org/neo4j/gds/similarity/filteredknn/FilteredKnnNeighborFilterFactory.java renamed to algo/src/main/java/org/neo4j/gds/similarity/knn/NeighbourConsumers.java

@@ -17,18 +17,16 @@
  * You should have received a copy of the GNU General Public License
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
-package org.neo4j.gds.similarity.filteredknn;
+package org.neo4j.gds.similarity.knn;
 
-public class FilteredKnnNeighborFilterFactory implements FilteredNeighborFilterFactory {
-
-    private final long nodeCount;
-
-    public FilteredKnnNeighborFilterFactory(long nodeCount) {
-        this.nodeCount = nodeCount;
-    }
+/**
+ * A holder for {@link org.neo4j.gds.similarity.knn.NeighbourConsumer}s. This instrument helps us extend KNN.
+ */
+public interface NeighbourConsumers {
+    /**
+     * A holder for sending data into the void, which is the default behaviour in regular KNN
+     */
+    NeighbourConsumers devNull = nodeId -> (element, priority) -> { /* do nothing */ };
 
-    @Override
-    public FilteredNeighborFilter create() {
-        return new FilteredKnnNeighborFilter(nodeCount);
-    }
+    NeighbourConsumer get(long nodeId);
 }

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

@@ -80,7 +80,8 @@ public long size() {
             k,
             k,
             Partition.of(0, nodeCount),
-            ProgressTracker.NULL_TRACKER
+            ProgressTracker.NULL_TRACKER,
+            NeighbourConsumers.devNull
         );
 
         generateRandomNeighbors.run();

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

@@ -36,7 +36,7 @@ class NeighborListTest {
     void shouldKeepMaxValuesOrderedByPriority() {
         long[] expected = {6L, 5L, 2L};
 
-        NeighborList queue = new NeighborList(3);
+        NeighborList queue = new NeighborList(3, NeighbourConsumer.devNull);
         SplittableRandom splittableRandom = new SplittableRandom();
 
         assertEquals(1, queue.add(0, Double.MIN_VALUE, splittableRandom, 0.0));
@@ -55,7 +55,7 @@ void shouldKeepMaxValuesOrderedByPriority() {
     void shouldLimitReturnWhenNotFull() {
         long[] expected = {6L, 5L, 4L};
 
-        NeighborList queue = new NeighborList(10);
+        NeighborList queue = new NeighborList(10, NeighbourConsumer.devNull);
         SplittableRandom splittableRandom = new SplittableRandom();
 
         assertEquals(1, queue.add(6, 6.0, splittableRandom, 0.0));
@@ -72,7 +72,7 @@ void insertEverything() {
         var elements = LongStream.range(0, nodeCount).boxed().collect(Collectors.toList());
         var rng = new SplittableRandom(1337L);
 
-        var queue = new NeighborList(nodeCount);
+        var queue = new NeighborList(nodeCount, NeighbourConsumer.devNull);
 
         elements.forEach(candidate -> queue.add(candidate, 1.0 / (1.0 + Math.abs(candidate - 2)), rng, 0.0));
 
@@ -82,7 +82,7 @@ void insertEverything() {
     @Test
     void insertEveryThingTake2() {
         List<Long> elements = List.of(0L, 2L);
-        var queue = new NeighborList(2);
+        var queue = new NeighborList(2, NeighbourConsumer.devNull);
         var rng = new SplittableRandom(1337L);
 
         elements.forEach(candidate -> queue.add(candidate, 1.0 / (1.0 + Math.abs(candidate - 1)), rng, 0.0));

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

@@ -50,7 +50,7 @@ void name(
 
         SplittableRandom rng = new SplittableRandom();
         allNeighbors.setAll(nodeId -> {
-            var neighbors = new NeighborList(k);
+            var neighbors = new NeighborList(k, NeighbourConsumer.devNull);
             LongStream.concat(
                 LongStream.range(nodeId + 1, nodeCount),
                 LongStream.range(0, nodeId)