Merge pull request #4242 from IoannisPanagiotas/bughuntingknn

Examining the fastrp+knn bug on filtered graphs.

Author: Mats-SX

algo/src/main/java/org/neo4j/gds/similarity/SimilarityGraphBuilder.java

@@ -71,16 +71,19 @@ public static MemoryEstimation memoryEstimation(int topK, int topN) {
     }
 
     private final NodeMapping nodeMapping;
+    private final NodeMapping rootNodeMapping;
     private final int concurrency;
     private final ExecutorService executorService;
     private final AllocationTracker allocationTracker;
 
     public SimilarityGraphBuilder(
         NodeMapping nodeMapping,
+        NodeMapping rootNodeMapping,
         int concurrency,
         ExecutorService executorService,
         AllocationTracker allocationTracker
     ) {
+        this.rootNodeMapping = rootNodeMapping;
         this.concurrency = concurrency;
         this.executorService = executorService;
         this.allocationTracker = allocationTracker;
@@ -100,7 +103,7 @@ public Graph build(Stream<SimilarityResult> stream) {
         ParallelUtil.parallelStreamConsume(stream, concurrency, relationshipsBuilder::addFromInternal);
 
         return GraphFactory.create(
-            nodeMapping,
+            rootNodeMapping,
             relationshipsBuilder.build(),
             allocationTracker
         );

algo/src/main/java/org/neo4j/gds/similarity/nodesim/NodeSimilarity.java

@@ -142,6 +142,7 @@ public SimilarityGraphResult computeToGraph() {
         } else {
             Stream<SimilarityResult> similarities = computeToStream();
             similarityGraph = new SimilarityGraphBuilder(
+                graph,
                 graph,
                 config.concurrency(),
                 executorService,

algo/src/test/java/org/neo4j/gds/similarity/nodesim/SimilarityGraphBuilderTest.java

@@ -20,8 +20,6 @@
 package org.neo4j.gds.similarity.nodesim;
 
 import org.junit.jupiter.api.Test;
-import org.neo4j.gds.similarity.SimilarityGraphBuilder;
-import org.neo4j.gds.similarity.SimilarityResult;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.core.concurrency.Pools;
 import org.neo4j.gds.core.huge.HugeGraph;
@@ -31,6 +29,8 @@
 import org.neo4j.gds.extension.GdlGraph;
 import org.neo4j.gds.extension.Inject;
 import org.neo4j.gds.extension.TestGraph;
+import org.neo4j.gds.similarity.SimilarityGraphBuilder;
+import org.neo4j.gds.similarity.SimilarityResult;
 
 import java.util.stream.Stream;
 
@@ -68,6 +68,7 @@ void testConstructionFromHugeGraph() {
         assertEquals(HugeGraph.class, unlabelledGraph.innerGraph().getClass());
 
         SimilarityGraphBuilder similarityGraphBuilder = new SimilarityGraphBuilder(
+            unlabelledGraph,
             unlabelledGraph,
             1,
             Pools.DEFAULT,
@@ -91,6 +92,7 @@ void testConstructionFromUnionGraph() {
         assertEquals(UnionGraph.class, graph.innerGraph().getClass());
 
         SimilarityGraphBuilder similarityGraphBuilder = new SimilarityGraphBuilder(
+            graph,
             graph,
             1,
             Pools.DEFAULT,

proc/similarity/src/main/java/org/neo4j/gds/similarity/knn/KnnMutateProc.java

@@ -118,6 +118,7 @@ protected Stream<SimilarityMutateResult> mutate(ComputationResult<Knn, Knn.Resul
             try (ProgressTimer ignored = ProgressTimer.start(mutateMillis::addAndGet)) {
                 similarityGraphResult = computeToGraph(
                     computationResult.graph(),
+                    computationResult.graphStore(),
                     algorithm.nodeCount(),
                     config.concurrency(),
                     Objects.requireNonNull(computationResult.result()),

proc/similarity/src/main/java/org/neo4j/gds/similarity/knn/KnnStatsProc.java

@@ -115,6 +115,7 @@ public Stream<SimilarityStatsResult> stats(AlgoBaseProc.ComputationResult<Knn, K
                 try (ProgressTimer ignored = resultBuilder.timePostProcessing()) {
                     SimilarityGraphResult similarityGraphResult = computeToGraph(
                         computationResult.graph(),
+                        computationResult.graphStore(),
                         algorithm.nodeCount(),
                         config.concurrency(),
                         result,

proc/similarity/src/main/java/org/neo4j/gds/similarity/knn/KnnWriteProc.java

@@ -21,6 +21,7 @@
 
 import org.neo4j.gds.AlgorithmFactory;
 import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.config.GraphCreateConfig;
 import org.neo4j.gds.core.CypherMapWrapper;
 import org.neo4j.gds.results.MemoryEstimateResult;
@@ -88,6 +89,7 @@ protected SimilarityGraphResult similarityGraphResult(ComputationResult<Knn, Knn
         KnnWriteConfig config = computationResult.config();
         return computeToGraph(
             computationResult.graph(),
+            computationResult.graphStore(),
             algorithm.nodeCount(),
             config.concurrency(),
             Objects.requireNonNull(computationResult.result()),
@@ -97,13 +99,15 @@ protected SimilarityGraphResult similarityGraphResult(ComputationResult<Knn, Knn
 
     static SimilarityGraphResult computeToGraph(
         Graph graph,
+        GraphStore graphStore,
         long nodeCount,
         int concurrency,
         Knn.Result result,
         KnnContext context
     ) {
         Graph similarityGraph = new SimilarityGraphBuilder(
             graph,
+            graphStore.nodes(),
             concurrency,
             context.executor(),
             context.allocationTracker()

proc/similarity/src/test/java/org/neo4j/gds/similarity/knn/KnnMutateProcTest.java

@@ -28,16 +28,21 @@
 import org.neo4j.gds.Orientation;
 import org.neo4j.gds.StoreLoaderBuilder;
 import org.neo4j.gds.api.DefaultValue;
+import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.nodeproperties.ValueType;
 import org.neo4j.gds.core.CypherMapWrapper;
+import org.neo4j.gds.core.loading.GraphStoreCatalog;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.lessThan;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.neo4j.gds.TestSupport.assertGraphEquals;
+import static org.neo4j.gds.TestSupport.fromGdl;
 
 class KnnMutateProcTest extends KnnProcTest<KnnMutateConfig>
     implements MutateRelationshipWithPropertyTest<Knn, KnnMutateConfig, Knn.Result> {
@@ -149,7 +154,7 @@ void shouldMutateResults() {
     @Override
     @Test
     @Disabled("This test does not work for KNN")
-    public void testGraphMutationOnFilteredGraph() { }
+    public void testGraphMutationOnFilteredGraph() {}
 
     @Test
     void shouldMutateUniqueRelationships() {
@@ -200,4 +205,50 @@ public void setupStoreLoader(StoreLoaderBuilder storeLoaderBuilder, Map<String,
             );
         }
     }
+
+    @Test
+    void shouldMutateWithFilteredNodes() {
+        String nodeCreateQuery =
+            "CREATE " +
+            "  (alice:Person {age: 24})" +
+            " ,(carol:Person {age: 24})" +
+            " ,(eve:Person {age: 67})" +
+            " ,(dave:Foo {age: 48})" +
+            " ,(bob:Foo {age: 48})";
+
+        runQuery(nodeCreateQuery);
+
+        String createQuery = GdsCypher.call()
+            .withNodeLabel("Person")
+            .withNodeLabel("Foo")
+            .withNodeProperty("age")
+            .withAnyRelationshipType()
+            .graphCreate("graph")
+            .yields();
+        runQuery(createQuery);
+
+        String relationshipType = "SIMILAR";
+        String relationshipProperty = "score";
+
+        String algoQuery = GdsCypher.call()
+            .explicitCreation("graph")
+            .algo("gds.beta.knn")
+            .mutateMode()
+            .addParameter("nodeLabels", List.of("Foo"))
+            .addParameter("nodeWeightProperty", "age")
+            .addParameter("mutateRelationshipType", relationshipType)
+            .addParameter("mutateProperty", relationshipProperty).yields();
+        runQuery(algoQuery);
+
+        Graph mutatedGraph = GraphStoreCatalog.get(getUsername(), db.databaseId(), "graph").graphStore().getUnion();
+
+        assertGraphEquals(
+            fromGdl(
+                nodeCreateQuery +
+                "(dave)-[{score: 1.0}]->(bob)" +
+                "(bob)-[{score: 1.0}]->(dave)"
+            ),
+            mutatedGraph
+        );
+    }
 }

proc/similarity/src/test/java/org/neo4j/gds/similarity/knn/KnnStreamProcTest.java

@@ -21,6 +21,7 @@
 
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.AlgoBaseProc;
+import org.neo4j.gds.GdsCypher;
 import org.neo4j.gds.core.CypherMapWrapper;
 import org.neo4j.gds.similarity.SimilarityResult;
 
@@ -62,4 +63,38 @@ void shouldStreamResults() {
             Map.of("node1", 2L, "node2", 1L, "similarity", 0.25)
         ));
     }
+
+    @Test
+    void shouldStreamWithFilteredNodes() {
+        String nodeCreateQuery =
+            "CREATE " +
+            "  (alice:Person {age: 24})" +
+            " ,(carol:Person {age: 24})" +
+            " ,(eve:Person {age: 67})" +
+            " ,(dave:Foo {age: 48})" +
+            " ,(bob:Foo {age: 48})";
+
+        runQuery(nodeCreateQuery);
+
+        String createQuery = GdsCypher.call()
+            .withNodeLabel("Person")
+            .withNodeLabel("Foo")
+            .withNodeProperty("age")
+            .withAnyRelationshipType()
+            .graphCreate("graph")
+            .yields();
+        runQuery(createQuery);
+
+        String algoQuery = GdsCypher.call()
+            .explicitCreation("graph")
+            .algo("gds.beta.knn")
+            .streamMode()
+            .addParameter("nodeLabels", List.of("Foo"))
+            .addParameter("nodeWeightProperty", "age")
+            .yields("node1", "node2", "similarity");
+        assertCypherResult(algoQuery, List.of(
+            Map.of("node1", 6L, "node2", 7L, "similarity", 1.0),
+            Map.of("node1", 7L, "node2", 6L, "similarity", 1.0)
+        ));
+    }
 }

proc/similarity/src/test/java/org/neo4j/gds/similarity/knn/KnnWriteProcTest.java

@@ -30,6 +30,8 @@
 import org.neo4j.gds.StoreLoaderBuilder;
 import org.neo4j.gds.WriteRelationshipWithPropertyTest;
 import org.neo4j.gds.api.DefaultValue;
+import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.core.Aggregation;
 import org.neo4j.gds.core.CypherMapWrapper;
 import org.neo4j.gds.core.loading.GraphStoreCatalog;
 import org.neo4j.gds.core.utils.progress.GlobalTaskStore;
@@ -38,6 +40,7 @@
 import org.neo4j.gds.test.config.ConcurrencyConfigProcTest;
 
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Stream;
@@ -202,7 +205,10 @@ void testProgressTracking() {
 
             var taskStore = new GlobalTaskStore();
 
-            pathProc.taskRegistryFactory = () -> new NonReleasingTaskRegistry(new TaskRegistry(getUsername(), taskStore));
+            pathProc.taskRegistryFactory = () -> new NonReleasingTaskRegistry(new TaskRegistry(
+                getUsername(),
+                taskStore
+            ));
 
             pathProc.write("undirectedGraph", createMinimalConfig(CypherMapWrapper.empty()).toMap());
 
@@ -213,6 +219,58 @@ void testProgressTracking() {
         });
     }
 
+    @Test
+    void shouldWriteWithFilteredNodes() {
+        runQuery("CREATE (alice:Person {name: 'Alice', age: 24})" +
+                 "CREATE (carol:Person {name: 'Carol', age: 24})" +
+                 "CREATE (eve:Person {name: 'Eve', age: 67})" +
+                 "CREATE (dave:Foo {name: 'Dave', age: 48})" +
+                 "CREATE (bob:Foo {name: 'Bob', age: 48})");
+
+        String createQuery = GdsCypher.call()
+            .withNodeLabel("Person")
+            .withNodeLabel("Foo")
+            .withNodeProperty("age")
+            .withAnyRelationshipType()
+            .graphCreate("graph")
+            .yields();
+        runQuery(createQuery);
+
+        String relationshipType = "SIMILAR";
+        String relationshipProperty = "score";
+
+        String algoQuery = GdsCypher.call()
+            .explicitCreation("graph")
+            .algo("gds.beta.knn")
+            .writeMode()
+            .addParameter("nodeLabels", List.of("Foo"))
+            .addParameter("nodeWeightProperty", "age")
+            .addParameter("writeRelationshipType", relationshipType)
+            .addParameter("writeProperty", relationshipProperty).yields();
+        runQuery(algoQuery);
+
+        Graph knnGraph = new StoreLoaderBuilder()
+            .api(db)
+            .addNodeLabel("Person")
+            .addNodeLabel("Foo")
+            .addRelationshipType(relationshipType)
+            .addRelationshipProperty(relationshipProperty, relationshipProperty, DefaultValue.DEFAULT, Aggregation.NONE)
+            .build()
+            .graph();
+
+        assertGraphEquals(
+            fromGdl("(alice:Person)" +
+                    "(carol:Person)" +
+                    "(eve:Person)" +
+                    "(dave:Foo)" +
+                    "(bob:Foo)" +
+                    "(dave)-[{score: 1.0}]->(bob)" +
+                    "(bob)-[{score: 1.0}]->(dave)"
+            ),
+            knnGraph
+        );
+    }
+
     @Override
     public String writeRelationshipType() {
         return "KNN_REL";