Implement relationshipSchema and correct relationshipCount for result graphs

Author: FlorentinD

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

@@ -19,10 +19,23 @@
  */
 package org.neo4j.gds.similarity.nodesim;
 
+import org.neo4j.gds.RelationshipType;
+import org.neo4j.gds.api.DefaultValue;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.GraphAdapter;
+import org.neo4j.gds.api.PropertyState;
+import org.neo4j.gds.api.nodeproperties.ValueType;
 import org.neo4j.gds.api.properties.relationships.RelationshipConsumer;
 import org.neo4j.gds.api.properties.relationships.RelationshipWithPropertyConsumer;
+import org.neo4j.gds.api.schema.Direction;
+import org.neo4j.gds.api.schema.GraphSchema;
+import org.neo4j.gds.api.schema.ImmutableMutableGraphSchema;
+import org.neo4j.gds.api.schema.ImmutableRelationshipPropertySchema;
+import org.neo4j.gds.api.schema.MutableRelationshipSchema;
+import org.neo4j.gds.api.schema.MutableRelationshipSchemaEntry;
+import org.neo4j.gds.core.Aggregation;
+
+import java.util.Map;
 
 public class TopKGraph extends GraphAdapter {
 
@@ -39,6 +52,44 @@ public int degree(long nodeId) {
         return topKList != null ? topKList.size() : 0;
     }
 
+    @Override
+    public boolean hasRelationshipProperty() {
+        return true;
+    }
+
+    @Override
+    public GraphSchema schema() {
+        var relationshipSchema = new MutableRelationshipSchema(
+            Map.of(
+                RelationshipType.ALL_RELATIONSHIPS,
+                new MutableRelationshipSchemaEntry(
+                    RelationshipType.ALL_RELATIONSHIPS,
+                    Direction.DIRECTED,
+                    Map.of(
+                        "similarity",
+                        ImmutableRelationshipPropertySchema.of(
+                            "similarity",
+                            ValueType.DOUBLE,
+                            DefaultValue.DEFAULT,
+                            PropertyState.TRANSIENT,
+                            Aggregation.NONE
+                        )
+                    )
+                )
+            )
+        );
+
+        return ImmutableMutableGraphSchema.builder()
+            .from(graph.schema())
+            .relationshipSchema(relationshipSchema)
+            .build();
+    }
+
+    @Override
+    public boolean isMultiGraph() {
+        return super.isMultiGraph();
+    }
+
     @Override
     public long relationshipCount() {
         return topKMap.similarityPairCount();

algo/src/main/java/org/neo4j/gds/spanningtree/SpanningGraph.java

@@ -19,12 +19,24 @@
  */
 package org.neo4j.gds.spanningtree;
 
+import org.neo4j.gds.RelationshipType;
+import org.neo4j.gds.api.DefaultValue;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.GraphAdapter;
+import org.neo4j.gds.api.PropertyState;
+import org.neo4j.gds.api.nodeproperties.ValueType;
 import org.neo4j.gds.api.properties.relationships.RelationshipConsumer;
 import org.neo4j.gds.api.properties.relationships.RelationshipWithPropertyConsumer;
+import org.neo4j.gds.api.schema.Direction;
+import org.neo4j.gds.api.schema.GraphSchema;
+import org.neo4j.gds.api.schema.ImmutableMutableGraphSchema;
+import org.neo4j.gds.api.schema.ImmutableRelationshipPropertySchema;
+import org.neo4j.gds.api.schema.MutableRelationshipSchema;
+import org.neo4j.gds.api.schema.MutableRelationshipSchemaEntry;
+import org.neo4j.gds.core.Aggregation;
 
 import java.util.Arrays;
+import java.util.Map;
 
 public class SpanningGraph extends GraphAdapter {
 
@@ -44,6 +56,50 @@ public int degree(long nodeId) {
         }
     }
 
+    @Override
+    public long relationshipCount() {
+        // parent -> node for each node (root -> root included)
+        return spanningTree.effectiveNodeCount();
+    }
+
+    @Override
+    public boolean hasRelationshipProperty() {
+        return true;
+    }
+
+    @Override
+    public GraphSchema schema() {
+        var relationshipSchema = new MutableRelationshipSchema(
+            Map.of(
+                RelationshipType.ALL_RELATIONSHIPS,
+                new MutableRelationshipSchemaEntry(
+                    RelationshipType.ALL_RELATIONSHIPS,
+                    Direction.DIRECTED,
+                    Map.of(
+                        "cost",
+                        ImmutableRelationshipPropertySchema.of(
+                            "cost",
+                            ValueType.DOUBLE,
+                            DefaultValue.DEFAULT,
+                            PropertyState.TRANSIENT,
+                            Aggregation.NONE
+                        )
+                    )
+                )
+            )
+        );
+
+        return ImmutableMutableGraphSchema.builder()
+            .from(graph.schema())
+            .relationshipSchema(relationshipSchema)
+            .build();
+    }
+
+    @Override
+    public boolean isMultiGraph() {
+        return false;
+    }
+
     @Override
     public void forEachRelationship(long nodeId, RelationshipConsumer consumer) {
         forEachRelationship(