Fix to undirected for relationships with a single property

Author: DarthMax

algo/src/main/java/org/neo4j/gds/beta/undirected/ToUndirected.java

@@ -19,19 +19,24 @@
  */
 package org.neo4j.gds.beta.undirected;
 
+import org.jetbrains.annotations.NotNull;
 import org.neo4j.gds.Algorithm;
 import org.neo4j.gds.Orientation;
 import org.neo4j.gds.RelationshipType;
 import org.neo4j.gds.api.CompositeRelationshipIterator;
+import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.api.PropertyState;
+import org.neo4j.gds.api.RelationshipIterator;
 import org.neo4j.gds.api.RelationshipProperty;
 import org.neo4j.gds.api.RelationshipPropertyStore;
 import org.neo4j.gds.api.schema.Direction;
 import org.neo4j.gds.api.schema.PropertySchema;
+import org.neo4j.gds.api.schema.RelationshipPropertySchema;
 import org.neo4j.gds.core.concurrency.RunWithConcurrency;
 import org.neo4j.gds.core.loading.SingleTypeRelationshipImportResult;
 import org.neo4j.gds.core.loading.construction.GraphFactory;
+import org.neo4j.gds.core.loading.construction.RelationshipsAndDirection;
 import org.neo4j.gds.core.loading.construction.RelationshipsBuilder;
 import org.neo4j.gds.core.loading.construction.RelationshipsBuilderBuilder;
 import org.neo4j.gds.core.utils.partition.DegreePartition;
@@ -42,6 +47,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -75,30 +81,9 @@ public SingleTypeRelationshipImportResult compute() {
             .propertySchemasFor(fromRelationshipType);
         var propertyKeys = propertySchemas.stream().map(PropertySchema::key).collect(Collectors.toList());
 
-        RelationshipsBuilderBuilder relationshipsBuilderBuilder = GraphFactory.initRelationshipsBuilder()
-            .concurrency(config.concurrency())
-            .nodes(graphStore.nodes())
-            .executorService(executorService)
-            .orientation(Orientation.UNDIRECTED)
-            .validateRelationships(false);
-
-        propertySchemas.forEach(propertySchema ->
-            relationshipsBuilderBuilder.addPropertyConfig(propertySchema.aggregation(), propertySchema.defaultValue())
-        );
+        var relationshipsBuilder = initializeRelationshipsBuilder(propertySchemas);
 
-        RelationshipsBuilder relationshipsBuilder = relationshipsBuilderBuilder.build();
-
-        CompositeRelationshipIterator relationshipIterator = graphStore.getCompositeRelationshipIterator(
-            fromRelationshipType,
-            propertyKeys
-        );
-
-        List<ToUndirectedTask> tasks = PartitionUtils.degreePartition(
-            graphStore.getGraph(fromRelationshipType),
-            config.concurrency(),
-            (partition) -> new ToUndirectedTask(relationshipsBuilder, relationshipIterator.concurrentCopy(), partition, progressTracker),
-            Optional.empty()
-        );
+        var tasks = createTasks(fromRelationshipType, propertyKeys, relationshipsBuilder);
 
         progressTracker.beginSubTask();
 
@@ -117,6 +102,38 @@ public SingleTypeRelationshipImportResult compute() {
         var relationships = relationshipsBuilder.buildAll();
         progressTracker.endSubTask();
 
+        SingleTypeRelationshipImportResult result = createResult(
+            propertySchemas,
+            propertyKeys,
+            relationships
+        );
+
+        progressTracker.endSubTask();
+
+        return result;
+    }
+
+    @NotNull
+    private RelationshipsBuilder initializeRelationshipsBuilder(List<RelationshipPropertySchema> propertySchemas) {
+        RelationshipsBuilderBuilder relationshipsBuilderBuilder = GraphFactory.initRelationshipsBuilder()
+            .concurrency(config.concurrency())
+            .nodes(graphStore.nodes())
+            .executorService(executorService)
+            .orientation(Orientation.UNDIRECTED)
+            .validateRelationships(false);
+
+        propertySchemas.forEach(propertySchema ->
+            relationshipsBuilderBuilder.addPropertyConfig(propertySchema.aggregation(), propertySchema.defaultValue())
+        );
+
+        return relationshipsBuilderBuilder.build();
+    }
+
+    private static SingleTypeRelationshipImportResult createResult(
+        List<RelationshipPropertySchema> propertySchemas,
+        List<String> propertyKeys,
+        List<RelationshipsAndDirection> relationships
+    ) {
         var topology = relationships.get(0).relationships().topology();
         var propertyValues = IntStream.range(0, propertyKeys.size())
             .boxed()
@@ -136,27 +153,96 @@ public SingleTypeRelationshipImportResult compute() {
             .putAllRelationshipProperties(propertyValues)
             .build();
 
-        progressTracker.endSubTask();
         return SingleTypeRelationshipImportResult.builder()
             .topology(topology)
             .properties(propertyStore)
             .direction(Direction.UNDIRECTED)
             .build();
     }
 
+    @NotNull
+    private List<Runnable> createTasks(
+        RelationshipType fromRelationshipType,
+        List<String> propertyKeys,
+        RelationshipsBuilder relationshipsBuilder
+    ) {
+        Function<DegreePartition, Runnable> taskCreator;
+        if (propertyKeys.size() == 1) {
+            Graph graph = graphStore.getGraph(fromRelationshipType, Optional.of(propertyKeys.get(0)));
+
+            taskCreator = (partition) -> new ToUndirectedTaskWithSingleProperty(
+                relationshipsBuilder,
+                graph.concurrentCopy(),
+                partition,
+                progressTracker
+            );
+        }
+        else {
+            CompositeRelationshipIterator relationshipIterator = graphStore.getCompositeRelationshipIterator(
+                fromRelationshipType,
+                propertyKeys
+            );
+
+            taskCreator = (partition) -> new ToUndirectedTaskWithMultipleProperties(
+                relationshipsBuilder,
+                relationshipIterator.concurrentCopy(),
+                partition,
+                progressTracker
+            );
+        }
+
+        return PartitionUtils.degreePartition(
+            graphStore.getGraph(fromRelationshipType),
+            config.concurrency(),
+            taskCreator,
+            Optional.empty()
+        );
+    }
+
     @Override
     public void release() {
 
     }
 
-    private static final class ToUndirectedTask implements Runnable {
+    private static final class ToUndirectedTaskWithSingleProperty implements Runnable {
+
+        private final RelationshipsBuilder relationshipsBuilder;
+        private final RelationshipIterator relationshipIterator;
+        private final DegreePartition partition;
+        private final ProgressTracker progressTracker;
+
+        private ToUndirectedTaskWithSingleProperty(
+            RelationshipsBuilder relationshipsBuilder,
+            RelationshipIterator relationshipIterator,
+            DegreePartition partition,
+            ProgressTracker progressTracker
+        ) {
+            this.relationshipsBuilder = relationshipsBuilder;
+            this.relationshipIterator = relationshipIterator;
+            this.partition = partition;
+            this.progressTracker = progressTracker;
+        }
+
+        @Override
+        public void run() {
+            for (long i = partition.startNode(); i < partition.startNode() + partition.nodeCount(); i++) {
+                relationshipIterator.forEachRelationship(i, 0.0D, (source, target, property) -> {
+                    relationshipsBuilder.addFromInternal(target, source, property);
+                    return true;
+                });
+                progressTracker.logProgress();
+            }
+        }
+    }
+
+    private static final class ToUndirectedTaskWithMultipleProperties implements Runnable {
 
         private final RelationshipsBuilder relationshipsBuilder;
         private final CompositeRelationshipIterator relationshipIterator;
         private final DegreePartition partition;
         private final ProgressTracker progressTracker;
 
-        private ToUndirectedTask(
+        private ToUndirectedTaskWithMultipleProperties(
             RelationshipsBuilder relationshipsBuilder,
             CompositeRelationshipIterator relationshipIterator,
             DegreePartition partition,
@@ -172,7 +258,7 @@ private ToUndirectedTask(
         public void run() {
             for (long i = partition.startNode(); i < partition.startNode() + partition.nodeCount(); i++) {
                 relationshipIterator.forEachRelationship(i, (source, target, properties) -> {
-                    relationshipsBuilder.add(target, source, properties);
+                    relationshipsBuilder.addFromInternal(target, source, properties);
                     return true;
                 });
                 progressTracker.logProgress();

proc/beta/src/test/java/org/neo4j/gds/beta/undirected/ToUndirectedProcTest.java

@@ -32,6 +32,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 
 import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -48,9 +49,9 @@ class ToUndirectedProcTest extends BaseProcTest {
                                     "(a:A) " +
                                     ",(b:B) " +
                                     ",(c:C) " +
-                                    ",(a)-[:REL]->(b)" +
-                                    ",(a)-[:REL]->(c)" +
-                                    ",(b)-[:REL]->(c)";
+                                    ",(a)-[:REL {prop1: 1.0}]->(b)" +
+                                    ",(a)-[:REL {prop1: 2.0}]->(c)" +
+                                    ",(b)-[:REL {prop1: 3.0}]->(c)";
 
     @BeforeEach
     void setup() throws Exception {
@@ -62,6 +63,7 @@ void setup() throws Exception {
             .withNodeLabel("B")
             .withNodeLabel("C")
             .withRelationshipType("REL")
+            .withRelationshipProperty("prop1")
             .yields()
         );
     }
@@ -80,8 +82,11 @@ void convertToUndirected() {
         ));
 
         var gs = GraphStoreCatalog.get("", "neo4j", "graph");
-        var graph = gs.graphStore().getGraph(RelationshipType.of("REL2"));
-        assertGraphEquals(fromGdl(DB.replace("REL", "REL2"), Orientation.UNDIRECTED), graph);
+        var graph = gs.graphStore().getGraph(RelationshipType.of("REL2"), Optional.of("prop1"));
+        assertGraphEquals(
+            fromGdl(DB.replace("REL", "REL2"), Orientation.UNDIRECTED),
+            graph
+        );
     }
 
     @Test