Use predicate type for source nodes also

Author: jjaderberg

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

@@ -31,7 +31,6 @@
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.similarity.knn.metrics.SimilarityComputer;
 
-import java.util.List;
 import java.util.Optional;
 import java.util.SplittableRandom;
 import java.util.concurrent.ExecutorService;
@@ -57,8 +56,8 @@ public class FilteredKnn extends Algorithm<FilteredKnnResult> {
     private final int sampledK;
     private final double sampleRate;
     private final double similarityCutoff;
-    private final List<Long> sourceNodes;
-    private final TargetNodePredicate targetNodePredicate;
+    private final NodeFilter sourceNodeFilter;
+    private final NodeFilter targetNodeFilter;
     private final int topK;
 
     private long nodePairsConsidered;
@@ -79,7 +78,8 @@ public static FilteredKnn create(
         var splittableRandom = getSplittableRandom(config.randomSeed());
         var sourceNodes = config.sourceNodeFilter().stream().map(graph::toMappedNodeId).collect(Collectors.toList());
         var targetNodes = config.targetNodeFilter().stream().map(graph::toMappedNodeId).collect(Collectors.toList());
-        var targetNodePredicate = new TargetNodePredicate(targetNodes);
+        var sourceNodeFilter = new NodeFilter(sourceNodes);
+        var targetNodeFilter = new NodeFilter(targetNodes);
         var samplerSupplier = samplerSupplier(graph, config);
         return new FilteredKnn(
             context.progressTracker(),
@@ -98,8 +98,8 @@ public static FilteredKnn create(
             config.sampledK(graph.nodeCount()),
             config.sampleRate(),
             config.similarityCutoff(),
-            sourceNodes,
-            targetNodePredicate,
+            sourceNodeFilter,
+            targetNodeFilter,
             config.topK()
         );
     }
@@ -142,8 +142,8 @@ private FilteredKnn(
         int sampledK,
         double sampleRate,
         double similarityCutoff,
-        List<Long> sourceNodes,
-        TargetNodePredicate targetNodePredicate,
+        NodeFilter sourceNodeFilter,
+        NodeFilter targetNodeFilter,
         int topK
 
     ) {
@@ -163,8 +163,8 @@ private FilteredKnn(
         this.sampledK = sampledK;
         this.sampleRate = sampleRate;
         this.similarityCutoff = similarityCutoff;
-        this.sourceNodes = sourceNodes;
-        this.targetNodePredicate = targetNodePredicate;
+        this.sourceNodeFilter = sourceNodeFilter;
+        this.targetNodeFilter = targetNodeFilter;
         this.topK = topK;
     }
 
@@ -223,7 +223,7 @@ public FilteredKnnResult compute() {
             this.progressTracker.endSubTask();
 
             this.progressTracker.endSubTask();
-            return ImmutableFilteredKnnResult.of(neighbors, iteration, didConverge, this.nodePairsConsidered, this.sourceNodes);
+            return ImmutableFilteredKnnResult.of(neighbors, iteration, didConverge, this.nodePairsConsidered, this.sourceNodeFilter);
         }
     }
 

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

@@ -42,7 +42,7 @@ public abstract class FilteredKnnResult {
 
     public abstract long nodePairsConsidered();
 
-    public abstract List<Long> sourceNodes();
+    public abstract NodeFilter sourceNodeFilter();
 
     public LongStream neighborsOf(long nodeId) {
         return neighborList().get(nodeId).elements().map(FilteredNeighborList::clearCheckedFlag);
@@ -55,7 +55,7 @@ public Stream<SimilarityResult> streamSimilarityResult() {
         return Stream
             .iterate(neighborList.initCursor(neighborList.newCursor()), HugeCursor::next, UnaryOperator.identity())
             .flatMap(cursor -> IntStream.range(cursor.offset, cursor.limit)
-//                .filter(index -> sourceNodes().contains(index + cursor.base))
+                .filter(index -> sourceNodeFilter().test(index + cursor.base))
                 .mapToObj(index -> cursor.array[index].similarityStream(index + cursor.base))
                 .flatMap(Function.identity())
             );
@@ -66,7 +66,7 @@ public long totalSimilarityPairs() {
         return Stream
             .iterate(neighborList.initCursor(neighborList.newCursor()), HugeCursor::next, UnaryOperator.identity())
             .flatMapToLong(cursor -> IntStream.range(cursor.offset, cursor.limit)
-//                .filter(index -> sourceNodes().contains(index + cursor.base))
+                .filter(index -> sourceNodeFilter().test(index + cursor.base))
                 .mapToLong(index -> cursor.array[index].size()))
             .sum();
     }
@@ -100,8 +100,8 @@ public long nodePairsConsidered() {
             }
 
             @Override
-            public List<Long> sourceNodes() {
-                return List.of();
+            public NodeFilter sourceNodeFilter() {
+                return new NodeFilter(List.of());
             }
 
             @Override

…ity/filteredknn/TargetNodePredicate.java …s/similarity/filteredknn/NodeFilter.javaalgo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodePredicate.java renamed to algo/src/main/java/org/neo4j/gds/similarity/filteredknn/NodeFilter.java algo/src/main/java/org/neo4j/gds/similarity/filteredknn/TargetNodePredicate.java renamed to algo/src/main/java/org/neo4j/gds/similarity/filteredknn/NodeFilter.java

@@ -21,10 +21,10 @@
 
 import java.util.List;
 
-public class TargetNodePredicate {
+public class NodeFilter {
     private final List<Long> nodeIds;
 
-    public TargetNodePredicate(List<Long> nodeIds) {this.nodeIds = nodeIds;}
+    public NodeFilter(List<Long> nodeIds) {this.nodeIds = nodeIds;}
 
     public boolean test(long nodeId) {
         return nodeIds.contains(nodeId);

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

@@ -26,7 +26,6 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.SplittableRandom;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.LongStream;
 import java.util.stream.Stream;
@@ -50,7 +49,7 @@ void should() {
             1,
             true,
             2,
-            List.of()
+            new NodeFilter(List.of())
         );
 
         var neighborLists = result.neighborList();