Add Memory Estimation for ToUndirected

Author: DarthMax

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

@@ -20,8 +20,14 @@
 package org.neo4j.gds.beta.undirected;
 
 import org.neo4j.gds.GraphStoreAlgorithmFactory;
+import org.neo4j.gds.RelationshipType;
 import org.neo4j.gds.api.GraphStore;
 import org.neo4j.gds.core.concurrency.Pools;
+import org.neo4j.gds.core.huge.CompressedAdjacencyList;
+import org.neo4j.gds.core.huge.UncompressedAdjacencyList;
+import org.neo4j.gds.core.utils.mem.MemoryEstimation;
+import org.neo4j.gds.core.utils.mem.MemoryEstimations;
+import org.neo4j.gds.core.utils.mem.MemoryRange;
 import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.core.utils.progress.tasks.Tasks;
@@ -50,4 +56,24 @@ public Task progressTask(GraphStore graphStore, ToUndirectedConfig config) {
             Tasks.leaf("Build undirected Adjacency list")
         );
     }
+
+    @Override
+    public MemoryEstimation memoryEstimation(ToUndirectedConfig configuration) {
+        RelationshipType relationshipType = RelationshipType.of(configuration.relationshipType());
+
+        var builder = MemoryEstimations.builder(ToUndirected.class)
+            .add("relationships", CompressedAdjacencyList.adjacencyListEstimation(relationshipType, true));
+
+        builder.perGraphDimension("properties", ((graphDimensions, integer) -> {
+            long max = graphDimensions.relationshipPropertyTokens().keySet().stream().mapToLong(__ ->
+                UncompressedAdjacencyList
+                    .adjacencyPropertiesEstimation(relationshipType, true)
+                    .estimate(graphDimensions, integer)
+                    .memoryUsage().max
+            ).sum();
+            return MemoryRange.of(0, max);
+        }));
+
+        return builder.build();
+    }
 }

executor/src/main/java/org/neo4j/gds/executor/GraphStoreFromCatalogLoader.java

@@ -32,6 +32,7 @@
 import org.neo4j.gds.core.loading.GraphStoreWithConfig;
 import org.neo4j.gds.core.loading.ImmutableCatalogRequest;
 
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
@@ -78,10 +79,17 @@ public GraphDimensions graphDimensions() {
         );
         long relCount = filteredGraph.relationshipCount();
 
+        var relationshipTypeTokens = new HashMap<String, Integer>();
+        var i = 0;
+        for (String key : graphStore.relationshipPropertyKeys()) {
+            relationshipTypeTokens.put(key, i++);
+        }
+
         return ImmutableGraphDimensions.builder()
             .nodeCount(filteredGraph.nodeCount())
             .relationshipCounts(filteredGraphRelationshipCounts(config, graphStore, filteredGraph))
             .relCountUpperBound(relCount)
+            .relationshipPropertyTokens(relationshipTypeTokens)
             .build();
     }
 

proc/beta/src/main/java/org/neo4j/gds/beta/undirected/ToUndirectedProc.java

@@ -21,8 +21,10 @@
 
 import org.neo4j.gds.BaseProc;
 import org.neo4j.gds.core.loading.SingleTypeRelationshipImportResult;
+import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.executor.ProcedureExecutor;
 import org.neo4j.gds.executor.ProcedureExecutorSpec;
+import org.neo4j.gds.results.MemoryEstimateResult;
 import org.neo4j.procedure.Description;
 import org.neo4j.procedure.Name;
 import org.neo4j.procedure.Procedure;
@@ -49,4 +51,20 @@ public Stream<ToUndirectedSpec.MutateResult> mutate(
             executionContext()
         ).compute(graphName, configuration, true, true);
     }
+
+    @Procedure(value = "gds.beta.graph.relationships.toUndirected.estimate", mode = READ)
+    @Description(ESTIMATE_DESCRIPTION)
+    public Stream<MemoryEstimateResult> estimate(
+        @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
+        @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
+    ) {
+        var mutateSpec = new ToUndirectedSpec();
+        var pipelineSpec = new ProcedureExecutorSpec<ToUndirected, SingleTypeRelationshipImportResult, ToUndirectedConfig>();
+
+        return new MemoryEstimationExecutor<>(
+            mutateSpec,
+            pipelineSpec,
+            executionContext()
+        ).computeEstimate(graphNameOrConfiguration, algoConfiguration);
+    }
 }

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

@@ -112,4 +112,21 @@ void shouldFailIfRelationshipTypeDoesNotExists() {
         );
         assertEquals(expectedMessage, throwable.getMessage());
     }
+
+    @Test
+    void memoryEstimation() {
+        String query = "CALL gds.beta.graph.relationships.toUndirected.estimate('graph', {relationshipType: 'REL', mutateRelationshipType: 'REL2'})";
+
+        assertCypherResult(query, List.of(Map.of(
+            "mapView", instanceOf(Map.class),
+            "treeView", instanceOf(String.class),
+            "bytesMax", 524632L,
+            "heapPercentageMin", 0.1D,
+            "nodeCount", 3L,
+            "requiredMemory", "[256 KiB ... 512 KiB]",
+            "bytesMin", 262360L,
+            "heapPercentageMax", 0.1D,
+            "relationshipCount", 3L
+        )));
+    }
 }