implementing a simple bounded priority queue for target node filtering

shortening names

Author: lassewesth

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

@@ -25,33 +25,46 @@
 import org.neo4j.gds.similarity.knn.Knn;
 import org.neo4j.gds.similarity.knn.KnnContext;
 
+<<<<<<< HEAD
+=======
+import java.util.List;
+
+/**
+ * Filtered KNN is the same as ordinary KNN, _but_ we allow users to regulate final output in two ways.
+ *
+ * Firstly, we enable source node filtering, meaning reported results are limited nodes from a certain set.
+ *
+ * Secondly, we enable target node filtering in the sense that every result will be from a certain set of nodes.
+ *
+ * In both cases the source or target node set can be actual specified nodes, or it could be all nodes with a label.
+ */
 public final class FilteredKnn extends Algorithm<FilteredKnnResult> {
     /**
      * This is KNN instrumented with neighbour consumers
      */
     private final Knn delegate;
 
-    private final TargetNodeFilteringNeighbourConsumers neighbourConsumers;
+    private final TargetNodeFiltering targetNodeFiltering;
     private final NodeFilter sourceNodeFilter;
 
     public static FilteredKnn create(Graph graph, FilteredKnnBaseConfig config, KnnContext context) {
         var targetNodeFilter = config.sourceNodeFilter().toNodeFilter(graph);
-        var neighbourConsumers = TargetNodeFilteringNeighbourConsumers.create(graph.nodeCount()/*, targetNodeFilter*/);
-        var knn = Knn.createWithDefaultsAndInstrumentation(graph, config, context, neighbourConsumers);
+        var targetNodeFiltering = TargetNodeFiltering.create(graph.nodeCount(), config.boundedK(graph.nodeCount())/*, targetNodeFilter*/);
+        var knn = Knn.createWithDefaultsAndInstrumentation(graph, config, context, targetNodeFiltering);
         var sourceNodeFilter = config.sourceNodeFilter().toNodeFilter(graph);
 
-        return new FilteredKnn(context.progressTracker(), knn, neighbourConsumers, sourceNodeFilter);
+        return new FilteredKnn(context.progressTracker(), knn, targetNodeFiltering, sourceNodeFilter);
     }
 
     private FilteredKnn(
         ProgressTracker progressTracker,
         Knn delegate,
-        TargetNodeFilteringNeighbourConsumers neighbourConsumers,
+        TargetNodeFiltering targetNodeFiltering,
         NodeFilter sourceNodeFilter
     ) {
         super(progressTracker);
         this.delegate = delegate;
-        this.neighbourConsumers = neighbourConsumers;
+        this.targetNodeFiltering = targetNodeFiltering;
         this.sourceNodeFilter = sourceNodeFilter;
     }
 
@@ -60,7 +73,7 @@ public FilteredKnnResult compute() {
         Knn.Result result = delegate.compute();
 
         return ImmutableFilteredKnnResult.of(
-            neighbourConsumers,
+            targetNodeFiltering,
             result.ranIterations(),
             result.didConverge(),
             result.nodePairsConsidered(),

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

@@ -27,7 +27,7 @@
 
 @ValueClass
 public abstract class FilteredKnnResult {
-    abstract TargetNodeFilteringNeighbourConsumers neighbourConsumers();
+    abstract TargetNodeFiltering neighbourConsumers();
 
     public abstract int ranIterations();
 
@@ -38,7 +38,7 @@ public abstract class FilteredKnnResult {
     public abstract List<Long> sourceNodes();
 
     public Stream<SimilarityResult> similarityResultStream() {
-        TargetNodeFilteringNeighbourConsumers neighbourConsumers = neighbourConsumers();
+        TargetNodeFiltering neighbourConsumers = neighbourConsumers();
         List<Long> sourceNodes = sourceNodes();
 
         return neighbourConsumers.asSimilarityResultStream(sourceNodes::contains);

…argetNodeFilteringNeighbourConsumer.java …larity/filteredknn/TargetNodeFilter.javaalgo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFilteringNeighbourConsumer.java renamed to algo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFilter.java algo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFilteringNeighbourConsumer.java renamed to algo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFilter.java

@@ -19,18 +19,38 @@
  */
 package org.neo4j.gds.similarity.filteredknn;
 
+import org.apache.commons.lang3.tuple.Pair;
 import org.neo4j.gds.similarity.SimilarityResult;
 import org.neo4j.gds.similarity.knn.NeighbourConsumer;
 
+import java.util.Comparator;
+import java.util.TreeSet;
 import java.util.stream.Stream;
 
-public class TargetNodeFilteringNeighbourConsumer implements NeighbourConsumer {
+/**
+ * We sort results by score, descending.
+ *
+ * For now a simple bounded priority queue that does _not_ handle duplicates.
+ */
+public class TargetNodeFilter implements NeighbourConsumer {
+    private final TreeSet<Pair<Double, Long>> priorityQueue = new TreeSet<>(Comparator.reverseOrder());
+    private final int bound;
+
+    public TargetNodeFilter(int bound) {
+        this.bound = bound;
+    }
+
     @Override
     public void offer(long element, double priority) {
-        throw new UnsupportedOperationException("TODO");
+        priorityQueue.add(Pair.of(priority, element));
+
+        if (priorityQueue.size() > bound) priorityQueue.pollLast();
     }
 
+    /**
+     * As part of an instrumentation of KNN this is a handy utility.
+     */
     Stream<SimilarityResult> asSimilarityStream(long nodeId) {
-        throw new UnsupportedOperationException("TODO");
+        return priorityQueue.stream().map(p -> new SimilarityResult(nodeId, p.getRight(), p.getLeft()));
     }
 }

…rgetNodeFilteringNeighbourConsumers.java …ity/filteredknn/TargetNodeFiltering.javaalgo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFilteringNeighbourConsumers.java renamed to algo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFiltering.java algo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFilteringNeighbourConsumers.java renamed to algo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodeFiltering.java

@@ -31,24 +31,24 @@
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
-public class TargetNodeFilteringNeighbourConsumers implements NeighbourConsumers {
-    private final HugeObjectArray<TargetNodeFilteringNeighbourConsumer> neighbourConsumers;
+public final class TargetNodeFiltering implements NeighbourConsumers {
+    private final HugeObjectArray<TargetNodeFilter> neighbourConsumers;
 
-    public TargetNodeFilteringNeighbourConsumers(HugeObjectArray<TargetNodeFilteringNeighbourConsumer> neighbourConsumers) {
-        this.neighbourConsumers = neighbourConsumers;
-    }
-
-    static TargetNodeFilteringNeighbourConsumers create(long nodeCount) {
-        HugeObjectArray<TargetNodeFilteringNeighbourConsumer> neighbourConsumers = HugeObjectArray.newArray(
-            TargetNodeFilteringNeighbourConsumer.class,
+    static TargetNodeFiltering create(long nodeCount, int k) {
+        HugeObjectArray<TargetNodeFilter> neighbourConsumers = HugeObjectArray.newArray(
+            TargetNodeFilter.class,
             nodeCount
         );
 
         for (int i = 0; i < nodeCount; i++) {
-            neighbourConsumers.set(i, new TargetNodeFilteringNeighbourConsumer());
+            neighbourConsumers.set(i, new TargetNodeFilter(k));
         }
 
-        return new TargetNodeFilteringNeighbourConsumers(neighbourConsumers);
+        return new TargetNodeFiltering(neighbourConsumers);
+    }
+
+    private TargetNodeFiltering(HugeObjectArray<TargetNodeFilter> neighbourConsumers) {
+        this.neighbourConsumers = neighbourConsumers;
     }
 
     @Override

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

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) "Neo4j"
+ * Neo4j Sweden AB [http://neo4j.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * Neo4j is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * 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;
+
+import org.junit.jupiter.api.Test;
+import org.neo4j.gds.similarity.SimilarityResult;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class TargetNodeFilterTest {
+    @Test
+    void shouldPrioritiseTargetNodes() {
+        TargetNodeFilter consumer = new TargetNodeFilter(3);
+
+        consumer.offer(23, 3.14);
+        consumer.offer(42, 1.61);
+        consumer.offer(87, 2.71);
+
+        assertThat(consumer.asSimilarityStream(117)).containsExactly(
+            new SimilarityResult(117, 23, 3.14),
+            new SimilarityResult(117, 87, 2.71),
+            new SimilarityResult(117, 42, 1.61)
+        );
+    }
+
+    @Test
+    void shouldOnlyKeepTopK() {
+        TargetNodeFilter consumer = new TargetNodeFilter(2);
+
+        consumer.offer(23, 3.14);
+        consumer.offer(42, 1.61);
+        consumer.offer(87, 2.71);
+
+        assertThat(consumer.asSimilarityStream(117)).containsExactly(
+            new SimilarityResult(117, 23, 3.14),
+            new SimilarityResult(117, 87, 2.71)
+        );
+    }
+}