Add memory estimation for AllPairsShortestPaths stream

Author: brs96

algo/src/main/java/org/neo4j/gds/allshortestpaths/AllShortestPathsMemoryEstimateDefinition.java

@@ -0,0 +1,63 @@
+/*
+ * 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.gds.allshortestpaths;
+
+import org.neo4j.gds.mem.MemoryEstimateDefinition;
+import org.neo4j.gds.mem.MemoryEstimation;
+import org.neo4j.gds.mem.MemoryEstimations;
+import org.neo4j.gds.mem.Estimate;
+import org.neo4j.gds.msbfs.MSBFSMemoryEstimation;
+
+public final class AllShortestPathsMemoryEstimateDefinition implements MemoryEstimateDefinition {
+
+    private final boolean hasRelationshipWeightProperty;
+
+    public AllShortestPathsMemoryEstimateDefinition(boolean hasRelationshipWeightProperty) {
+        this.hasRelationshipWeightProperty = hasRelationshipWeightProperty;
+    }
+
+    @Override
+    public MemoryEstimation memoryEstimation() {
+        if (hasRelationshipWeightProperty) {
+            return weightedMemoryEstimation();
+        } else {
+            return unweightedMemoryEstimation();
+        }
+    }
+
+    private MemoryEstimation weightedMemoryEstimation() {
+        return MemoryEstimations.builder(WeightedAllShortestPaths.class)
+            .perThread("ShortestPathTask", shortestPathTaskMemoryEstimation())
+            .build();
+    }
+
+    private MemoryEstimation unweightedMemoryEstimation() {
+        return MemoryEstimations.builder(MSBFSAllShortestPaths.class)
+            .add("MSBFS", MSBFSMemoryEstimation.MSBFSWithANPStrategy(0))
+            .build();
+    }
+
+    private MemoryEstimation shortestPathTaskMemoryEstimation() {
+        return MemoryEstimations.builder(WeightedAllShortestPaths.ShortestPathTask.class)
+            .perNode("distance array", Estimate::sizeOfDoubleArray)
+            .add("priority queue", IntPriorityQueue.memoryEstimation())
+            .build();
+    }
+} 

algo/src/main/java/org/neo4j/gds/allshortestpaths/IntPriorityQueue.java

@@ -21,6 +21,9 @@
 import org.jetbrains.annotations.TestOnly;
 import org.neo4j.gds.collections.ArrayUtil;
 import org.neo4j.gds.collections.ha.HugeIntArray;
+import org.neo4j.gds.mem.MemoryEstimation;
+import org.neo4j.gds.mem.MemoryEstimations;
+import org.neo4j.gds.mem.Estimate;
 
 /**
  * A PriorityQueue specialized for ints that maintains a partial ordering of
@@ -44,6 +47,14 @@ public abstract class IntPriorityQueue {
     private final IntLongScatterMap mapElementToIndex;
     private long size = 0;
 
+    public static MemoryEstimation memoryEstimation() {
+        return MemoryEstimations.builder(IntPriorityQueue.class)
+            .perNode("heap", HugeIntArray::memoryEstimation)
+            .add("costs", MemoryEstimations.builder(IntDoubleScatterMap.class).build())
+            .add("element to index map", MemoryEstimations.builder(IntLongScatterMap.class).build())
+            .build();
+    }
+
     /**
      * Creates a new queue with the given capacity.
      * The queue dynamically grows to hold all elements.

algo/src/main/java/org/neo4j/gds/allshortestpaths/WeightedAllShortestPaths.java

@@ -109,7 +109,7 @@ public Stream<AllShortestPathsStreamResult> compute() {
      * Dijkstra Task. Takes one element of the counter at a time
      * and starts dijkstra on it.
      */
-    private final class ShortestPathTask implements Runnable {
+    public final class ShortestPathTask implements Runnable {
 
         private final IntPriorityQueue queue;
         private final double[] distance;

algo/src/test/java/org/neo4j/gds/allshortestpaths/AllShortestPathsMemoryEstimateDefinitionTest.java

@@ -0,0 +1,48 @@
+/*
+ * 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.gds.allshortestpaths;
+
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.neo4j.gds.assertions.MemoryEstimationAssert;
+import org.neo4j.gds.core.concurrency.Concurrency;
+
+import static org.neo4j.gds.assertions.MemoryEstimationAssert.assertThat;
+
+class AllShortestPathsMemoryEstimateDefinitionTest {
+
+    @ParameterizedTest
+    @CsvSource({
+        "10_000, 1, false, 280",
+        "10_000, 4, false, 544",
+        "500_000, 4, false, 27_200",
+        "10_000_000, 4, false, 544_000",
+        "10_000, 1, true, 1_120_456",
+        "10_000, 4, true, 1_120_720",
+        "500_000, 4, true, 56_000_720",
+        "10_000_000, 4, true, 1_120_000_720"
+    })
+    void testMemoryEstimation(long nodeCount, int concurrency, boolean weighted, long expectedMemory) {
+        var memoryEstimation = new AllShortestPathsMemoryEstimateDefinition(weighted).memoryEstimation();
+        assertThat(memoryEstimation)
+            .memoryRange(nodeCount, new Concurrency(concurrency))
+            .hasSameMinAndMaxEqualTo(expectedMemory);
+    }
+} 

applications/algorithms/path-finding/src/main/java/org/neo4j/gds/applications/algorithms/pathfinding/PathFindingAlgorithmsEstimationModeBusinessFacade.java

@@ -49,6 +49,8 @@
 import org.neo4j.gds.traversal.RandomWalkCountingVisitsMemoryEstimateDefinition;
 import org.neo4j.gds.traversal.RandomWalkMemoryEstimateDefinition;
 import org.neo4j.gds.traversal.RandomWalkMutateConfig;
+import org.neo4j.gds.allshortestpaths.AllShortestPathsMemoryEstimateDefinition;
+import org.neo4j.gds.allshortestpaths.AllShortestPathsConfig;
 
 /**
  * Here is the top level business facade for all your path finding memory estimation needs.
@@ -61,8 +63,14 @@ public PathFindingAlgorithmsEstimationModeBusinessFacade(AlgorithmEstimationTemp
         this.algorithmEstimationTemplate = algorithmEstimationTemplate;
     }
 
-    MemoryEstimation allShortestPaths() {
-        throw new MemoryEstimationNotImplementedException();
+    public MemoryEstimation allShortestPaths(AllShortestPathsConfig configuration) {
+        return new AllShortestPathsMemoryEstimateDefinition(configuration.hasRelationshipWeightProperty()).memoryEstimation();
+    }
+
+    public MemoryEstimateResult allShortestPaths(AllShortestPathsConfig configuration, Object graphNameOrConfiguration) {
+        var memoryEstimation = allShortestPaths(configuration);
+
+        return runEstimation(configuration, graphNameOrConfiguration, memoryEstimation);
     }
 
     public MemoryEstimateResult bellmanFord(AllShortestPathsBellmanFordBaseConfig configuration, Object graphNameOrConfiguration) {

open-packaging/src/test/java/org/neo4j/gds/OpenGdsProcedureSmokeTest.java

@@ -53,6 +53,7 @@ class OpenGdsProcedureSmokeTest extends BaseProcTest {
         "gds.graph.sample.cnarw.estimate",
 
         "gds.allShortestPaths.stream",
+        "gds.allShortestPaths.stream.estimate",
 
         "gds.articulationPoints.mutate",
         "gds.articulationPoints.mutate.estimate",
@@ -619,7 +620,7 @@ void countShouldMatch() {
         );
 
         // If you find yourself updating this count, please also update the count in SmokeTest.kt
-        int expectedCount = 459;
+        int expectedCount = 460;
         assertEquals(
             expectedCount,
             returnedRows,

proc/path-finding/src/main/java/org/neo4j/gds/paths/all/AllShortestPathsStreamProc.java

@@ -21,6 +21,7 @@
 
 import org.neo4j.gds.allshortestpaths.AllShortestPathsStreamResult;
 import org.neo4j.gds.procedures.GraphDataScienceProcedures;
+import org.neo4j.gds.applications.algorithms.machinery.MemoryEstimateResult;
 import org.neo4j.procedure.Context;
 import org.neo4j.procedure.Description;
 import org.neo4j.procedure.Internal;
@@ -31,6 +32,7 @@
 import java.util.stream.Stream;
 
 import static org.neo4j.gds.paths.all.Constants.ALL_PAIRS_SHORTEST_PATH_DESCRIPTION;
+import static org.neo4j.gds.procedures.ProcedureConstants.MEMORY_ESTIMATION_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 
 public class AllShortestPathsStreamProc {
@@ -46,6 +48,15 @@ public Stream<AllShortestPathsStreamResult> stream(
         return facade.algorithms().pathFinding().allShortestPathStream(graphName, configuration);
     }
 
+    @Procedure(name = "gds.allShortestPaths.stream.estimate", mode = READ)
+    @Description(MEMORY_ESTIMATION_DESCRIPTION)
+    public Stream<MemoryEstimateResult> estimate(
+        @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
+        @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
+    ) {
+        return facade.algorithms().pathFinding().allShortestPathStreamEstimate(graphNameOrConfiguration, algoConfiguration);
+    }
+
     @Procedure(name = "gds.alpha.allShortestPaths.stream", mode = READ, deprecatedBy = "gds.allShortestPaths.stream")
     @Description(ALL_PAIRS_SHORTEST_PATH_DESCRIPTION)
     @Internal

procedures/algorithms-facade/src/main/java/org/neo4j/gds/procedures/algorithms/pathfinding/LocalPathFindingProcedureFacade.java

@@ -265,6 +265,22 @@ public Stream<AllShortestPathsStreamResult> allShortestPathStream(
         );
     }
 
+    @Override
+    public Stream<MemoryEstimateResult> allShortestPathStreamEstimate(
+        Object graphNameOrConfiguration,
+        Map<String, Object> algorithmConfiguration
+    ) {
+        var parsedConfiguration = configurationParser.parseConfiguration(
+            algorithmConfiguration,
+            AllShortestPathsConfig::of
+        );
+
+        return Stream.of(estimationModeBusinessFacade.allShortestPaths(
+            parsedConfiguration,
+            graphNameOrConfiguration
+        ));
+    }
+
     @Override
     public Stream<BellmanFordStreamResult> bellmanFordStream(String graphName, Map<String, Object> configuration) {
         var routeRequested = procedureReturnColumns.contains("route");

procedures/facade-api/path-finding-facade-api/src/main/java/org/neo4j/gds/procedures/algorithms/pathfinding/PathFindingProcedureFacade.java

@@ -38,6 +38,11 @@ Stream<AllShortestPathsStreamResult> allShortestPathStream(
         Map<String, Object> configuration
     );
 
+    Stream<MemoryEstimateResult> allShortestPathStreamEstimate(
+        Object graphNameOrConfiguration,
+        Map<String, Object> algorithmConfiguration
+    );
+
     Stream<BellmanFordStreamResult> bellmanFordStream(String graphName, Map<String, Object> configuration);
 
     Stream<MemoryEstimateResult> bellmanFordStreamEstimate(