integrating target node filter

adding smoke test for target node filtering

Author: lassewesth

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

@@ -45,8 +45,8 @@ public final class FilteredKnn extends Algorithm<FilteredKnnResult> {
     private final NodeFilter sourceNodeFilter;
 
     public static FilteredKnn create(Graph graph, FilteredKnnBaseConfig config, KnnContext context) {
-        var targetNodeFilter = config.sourceNodeFilter().toNodeFilter(graph);
-        var targetNodeFiltering = TargetNodeFiltering.create(graph.nodeCount(), config.boundedK(graph.nodeCount())/*, targetNodeFilter*/);
+        var targetNodeFilter = config.targetNodeFilter().toNodeFilter(graph);
+        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);
 

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

@@ -25,6 +25,7 @@
 
 import java.util.Comparator;
 import java.util.TreeSet;
+import java.util.function.LongPredicate;
 import java.util.stream.Stream;
 
 /**
@@ -34,14 +35,19 @@
  */
 public class TargetNodeFilter implements NeighbourConsumer {
     private final TreeSet<Pair<Double, Long>> priorityQueue = new TreeSet<>(Comparator.reverseOrder());
+
+    private final LongPredicate predicate;
     private final int bound;
 
-    public TargetNodeFilter(int bound) {
+    public TargetNodeFilter(LongPredicate predicate, int bound) {
+        this.predicate = predicate;
         this.bound = bound;
     }
 
     @Override
     public void offer(long element, double priority) {
+        if (! predicate.test(element)) return;
+
         priorityQueue.add(Pair.of(priority, element));
 
         if (priorityQueue.size() > bound) priorityQueue.pollLast();

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

@@ -34,14 +34,14 @@
 public final class TargetNodeFiltering implements NeighbourConsumers {
     private final HugeObjectArray<TargetNodeFilter> neighbourConsumers;
 
-    static TargetNodeFiltering create(long nodeCount, int k) {
+    static TargetNodeFiltering create(long nodeCount, int k, LongPredicate targetNodePredicate) {
         HugeObjectArray<TargetNodeFilter> neighbourConsumers = HugeObjectArray.newArray(
             TargetNodeFilter.class,
             nodeCount
         );
 
         for (int i = 0; i < nodeCount; i++) {
-            neighbourConsumers.set(i, new TargetNodeFilter(k));
+            neighbourConsumers.set(i, new TargetNodeFilter(targetNodePredicate, k));
         }
 
         return new TargetNodeFiltering(neighbourConsumers);

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

@@ -172,4 +172,61 @@ void shouldOnlyProduceResultsForMultipleFilteredSourceNode() {
             ).isEqualTo(Set.of(filteredNode1, filteredNode2));
         }
     }
+
+    @Nested
+    class TargetNodeFiltering {
+        @GdlGraph
+        private static final String DB_CYPHER =
+            "CREATE" +
+            "  (a { knn: 1.2 } )" +
+            ", (b { knn: 1.1 } )" +
+            ", (c { knn: 2.1 } )" +
+            ", (d { knn: 3.1 } )" +
+            ", (e { knn: 4.1 } )";
+
+        @Test
+        void shouldOnlyProduceResultsForFilteredTargetNode() {
+            var targetNode = idFunction.of("a");
+            var config = FilteredKnnBaseConfigImpl.builder()
+                .nodeProperties(List.of("knn"))
+                .topK(3)
+                .randomJoins(0)
+                .maxIterations(1)
+                .randomSeed(20L)
+                .concurrency(1)
+                .targetNodeFilter(targetNode)
+                .build();
+            var knnContext = KnnContext.empty();
+            var knn = FilteredKnn.create(graph, config, knnContext);
+            var result = knn.compute();
+
+            assertThat(result.similarityResultStream()
+                .map(SimilarityResult::targetNodeId)
+                .collect(Collectors.toSet())
+            ).isEqualTo(Set.of(targetNode));
+        }
+
+        @Test
+        void shouldOnlyProduceResultsForFilteredTargetNodes() {
+            var targetNode1 = idFunction.of("a");
+            var targetNode2 = idFunction.of("b");
+            var config = FilteredKnnBaseConfigImpl.builder()
+                .nodeProperties("knn")
+                .topK(3)
+                .randomJoins(0)
+                .maxIterations(1)
+                .randomSeed(20L)
+                .concurrency(1)
+                .targetNodeFilter(List.of(targetNode1, targetNode2))
+                .build();
+            var knnContext = KnnContext.empty();
+            var knn = FilteredKnn.create(graph, config, knnContext);
+            var result = knn.compute();
+
+            assertThat(result.similarityResultStream()
+                .map(SimilarityResult::targetNodeId)
+                .collect(Collectors.toSet())
+            ).isEqualTo(Set.of(targetNode1, targetNode2));
+        }
+    }
 }

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

@@ -27,7 +27,7 @@
 class TargetNodeFilterTest {
     @Test
     void shouldPrioritiseTargetNodes() {
-        TargetNodeFilter consumer = new TargetNodeFilter(3);
+        TargetNodeFilter consumer = new TargetNodeFilter(l -> true, 3);
 
         consumer.offer(23, 3.14);
         consumer.offer(42, 1.61);
@@ -42,7 +42,7 @@ void shouldPrioritiseTargetNodes() {
 
     @Test
     void shouldOnlyKeepTopK() {
-        TargetNodeFilter consumer = new TargetNodeFilter(2);
+        TargetNodeFilter consumer = new TargetNodeFilter(l -> true, 2);
 
         consumer.offer(23, 3.14);
         consumer.offer(42, 1.61);
@@ -53,4 +53,15 @@ void shouldOnlyKeepTopK() {
             new SimilarityResult(117, 87, 2.71)
         );
     }
+
+    @Test
+    void shouldOnlyIncludeTargetNodes() {
+        TargetNodeFilter consumer = new TargetNodeFilter(l -> false, 3);
+
+        consumer.offer(23, 3.14);
+        consumer.offer(42, 1.61);
+        consumer.offer(87, 2.71);
+
+        assertThat(consumer.asSimilarityStream(117)).isEmpty();
+    }
 }