Fix AIIOB in Node Sim and KNN

Co-authored-by: Sören Reichardt <[email protected]>
Co-authored-by: Max Kießling <[email protected]>
Co-Authored-By: Paul Horn <[email protected]>
Co-Authored-By: Lasse Westh-Nielsen <[email protected]>

Author: 5 people

core/src/main/java/org/neo4j/graphalgo/core/huge/UnionGraph.java

@@ -75,7 +75,7 @@ public long nodeCount() {
 
     @Override
     public long rootNodeCount() {
-        return first.nodeCount();
+        return first.rootNodeCount();
     }
 
     @Override

core/src/test/java/org/neo4j/graphalgo/core/huge/UnionGraphTest.java

@@ -24,6 +24,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
+import org.neo4j.graphalgo.NodeLabel;
 import org.neo4j.graphalgo.Orientation;
 import org.neo4j.graphalgo.RelationshipType;
 import org.neo4j.graphalgo.api.Graph;
@@ -38,6 +39,7 @@
 
 import java.util.Collection;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Stream;
 
@@ -186,4 +188,24 @@ void overlappingRelationshipProperty() {
         assertThat(graph.relationshipProperty(a, b, -1)).isIn(3D, 5D);
     }
 
+    @Test
+    void filteredGraphRootNodeCountShouldAgreeWithUnfilteredGraph() {
+        // A graph that is both label filtered and has multiple relationship types
+        var graphStore = GdlFactory.of(
+            "CREATE " +
+            "  (a:A)" +
+            ", (b:B)" +
+            ", (a)-[:R1]->(b)" +
+            ", (b)-[:R2]->(a)"
+        ).build().graphStore();
+
+        var unionGraph = graphStore.getGraph(
+            List.of(NodeLabel.of("A")),
+            List.of(RelationshipType.of("R1"), RelationshipType.of("R2")),
+            Optional.empty()
+        );
+
+        assertThat(unionGraph.rootNodeCount()).isEqualTo(graphStore.nodeCount());
+    }
+
 }

proc/similarity/src/test/java/org/neo4j/graphalgo/similarity/knn/KnnFilteredGraphTest.java

@@ -0,0 +1,81 @@
+/*
+ * 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.graphalgo.similarity.knn;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.neo4j.graphalgo.BaseProcTest;
+import org.neo4j.graphalgo.catalog.GraphCreateProc;
+import org.neo4j.graphalgo.catalog.GraphStreamRelationshipPropertiesProc;
+import org.neo4j.graphalgo.extension.Neo4jGraph;
+
+import static org.assertj.core.api.Assertions.assertThatCode;
+
+class KnnFilteredGraphTest extends BaseProcTest {
+
+    @BeforeEach
+    void setUp() throws Exception {
+        registerProcedures(GraphCreateProc.class, KnnMutateProc.class, GraphStreamRelationshipPropertiesProc.class);
+    }
+
+    @Neo4jGraph
+    static String cypher =
+        "CREATE" +
+        "  (peter:Person {knn: 3.0})" +
+        ", (paul:Person {knn: 3.0})" +
+        ", (secretariat:Horse {knn: 2.0})" +
+        ", (ina_scott:Horse {knn: 3.0})" +
+        ", (mary:Person {knn: 3.0})" +
+        ", (:Person {knn: 3.0})" +
+        ", (lexington:Horse {knn: 3.0})" +
+        ", (:Person {knn: 3.0})" +
+        ", (peter)-[:LIKES]->(mary)" +
+        ", (paul)-[:KNOWS]->(peter)" +
+        ", (secretariat)-[:LIKES]->(ina_scott)" +
+        ", (lexington)-[:KNOWS]->(secretariat)";
+
+
+    @Test
+    void shouldComputeAndMutateOnFilteredGraph() {
+        var create = "CALL gds.graph.create(" +
+                         "'g'," +
+                         "['Person', 'Horse']," +
+                         "['LIKES', 'KNOWS']" +
+                         ", {" +
+                             "nodeProperties: ['knn']" +
+                         "}" +
+                     ")";
+        runQuery(create);
+        var knnQuery1 = "CALL gds.beta.knn.mutate('g', {" +
+                       "    nodeLabels: ['Horse']," +
+                       "    mutateRelationshipType: 'FOO1'," +
+                       "    mutateProperty: 'score'," +
+                       "    nodeWeightProperty: 'knn'" +
+                       "})";
+        var knnQuery2 = "CALL gds.beta.knn.mutate('g', {" +
+                        "    nodeLabels: ['Person']," +
+                        "    mutateRelationshipType: 'FOO2'," +
+                        "    mutateProperty: 'score'," +
+                        "    nodeWeightProperty: 'knn'" +
+                        "})";
+        assertThatCode(() -> runQuery(knnQuery1)).doesNotThrowAnyException();
+        assertThatCode(() -> runQuery(knnQuery2)).doesNotThrowAnyException();
+    }
+}

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

@@ -141,6 +141,7 @@ void shouldMutateResults() {
         });
     }
 
+
     @Override
     @Test
     @Disabled("This test does not work for KNN")