migrate last single pair shortest path estimate procedures to being r…

…eusable
neo4j · Jan 5, 2024 · e0f51c5 · e0f51c5
1 parent 904c2a1
commit e0f51c5
Show file tree

Hide file tree

Showing 9 changed files with 174 additions and 92 deletions.
diff --git a/...n/java/org/neo4j/gds/applications/algorithms/pathfinding/PathFindingAlgorithmsFacade.java b/...n/java/org/neo4j/gds/applications/algorithms/pathfinding/PathFindingAlgorithmsFacade.java
@@ -23,6 +23,7 @@
 import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
 import org.neo4j.gds.core.write.RelationshipStreamExporterBuilder;
 import org.neo4j.gds.logging.Log;
+import org.neo4j.gds.paths.SourceTargetShortestPathBaseConfig;
 import org.neo4j.gds.paths.astar.AStarMemoryEstimateDefinition;
 import org.neo4j.gds.paths.astar.config.ShortestPathAStarBaseConfig;
 import org.neo4j.gds.paths.astar.config.ShortestPathAStarMutateConfig;
@@ -36,6 +37,7 @@
 import org.neo4j.gds.paths.dijkstra.config.ShortestPathDijkstraStreamConfig;
 import org.neo4j.gds.paths.dijkstra.config.ShortestPathDijkstraWriteConfig;
 import org.neo4j.gds.paths.yens.YensMemoryEstimateDefinition;
+import org.neo4j.gds.paths.yens.config.ShortestPathYensBaseConfig;
 import org.neo4j.gds.paths.yens.config.ShortestPathYensMutateConfig;
 import org.neo4j.gds.paths.yens.config.ShortestPathYensStreamConfig;
 import org.neo4j.gds.paths.yens.config.ShortestPathYensWriteConfig;
@@ -160,6 +162,19 @@ public <RESULT> RESULT singlePairShortestPathAStarWrite(
         );
     }
 
+    public MemoryEstimateResult singlePairShortestPathDijkstraEstimate(
+        SourceTargetShortestPathBaseConfig configuration,
+        Object graphNameOrConfiguration
+    ) {
+        var memoryEstimation = new DijkstraMemoryEstimateDefinition().memoryEstimation(configuration);
+
+        return algorithmEstimationTemplate.estimate(
+            configuration,
+            graphNameOrConfiguration,
+            memoryEstimation
+        );
+    }
+
     public <RESULT> RESULT singlePairShortestPathDijkstraMutate(
         GraphName graphName,
         ShortestPathDijkstraMutateConfig configuration,
@@ -218,6 +233,19 @@ public <RESULT> RESULT singlePairShortestPathDijkstraWrite(
         );
     }
 
+    public MemoryEstimateResult singlePairShortestPathYensEstimate(
+        ShortestPathYensBaseConfig configuration,
+        Object graphNameOrConfiguration
+    ) {
+        var memoryEstimation = new YensMemoryEstimateDefinition().memoryEstimation(configuration);
+
+        return algorithmEstimationTemplate.estimate(
+            configuration,
+            graphNameOrConfiguration,
+            memoryEstimation
+        );
+    }
+
     public <RESULT> RESULT singlePairShortestPathYensMutate(
         GraphName graphName,
         ShortestPathYensMutateConfig configuration,

diff --git a/...th-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathAStarWriteProc.java b/...th-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathAStarWriteProc.java
@@ -19,10 +19,6 @@
  */
 package org.neo4j.gds.paths.sourcetarget;
 
-import org.neo4j.gds.BaseProc;
-import org.neo4j.gds.core.write.RelationshipStreamExporterBuilder;
-import org.neo4j.gds.executor.ExecutionContext;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.procedures.GraphDataScience;
 import org.neo4j.gds.results.MemoryEstimateResult;
 import org.neo4j.gds.results.StandardWriteRelationshipsResult;
@@ -34,17 +30,15 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.neo4j.gds.ProcedureConstants.ESTIMATE_DESCRIPTION;
 import static org.neo4j.gds.paths.sourcetarget.SinglePairShortestPathConstants.ASTAR_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 import static org.neo4j.procedure.Mode.WRITE;
 
-public class ShortestPathAStarWriteProc extends BaseProc {
+public class ShortestPathAStarWriteProc {
     @Context
     public GraphDataScience facade;
 
-    @Context
-    public RelationshipStreamExporterBuilder relationshipStreamExporterBuilder;
-
     @Procedure(name = "gds.shortestPath.astar.write", mode = WRITE)
     @Description(ASTAR_DESCRIPTION)
     public Stream<StandardWriteRelationshipsResult> write(
@@ -60,16 +54,9 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
         @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
     ) {
-        return new MemoryEstimationExecutor<>(
-            new ShortestPathAStarWriteSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphNameOrConfiguration, algoConfiguration);
+        return facade.pathFinding().singlePairShortestPathAStarWriteEstimate(
+            graphNameOrConfiguration,
+            algoConfiguration
+        );
     }
-
-    @Override
-    public ExecutionContext executionContext() {
-        return super.executionContext().withRelationshipStreamExporterBuilder(relationshipStreamExporterBuilder);
-    }
-
 }
diff --git a/...inding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathDijkstraMutateProc.java b/...inding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathDijkstraMutateProc.java
@@ -19,8 +19,6 @@
  */
 package org.neo4j.gds.paths.sourcetarget;
 
-import org.neo4j.gds.BaseProc;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.procedures.GraphDataScience;
 import org.neo4j.gds.procedures.pathfinding.PathFindingMutateResult;
 import org.neo4j.gds.results.MemoryEstimateResult;
@@ -32,10 +30,11 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.neo4j.gds.ProcedureConstants.ESTIMATE_DESCRIPTION;
 import static org.neo4j.gds.paths.sourcetarget.SinglePairShortestPathConstants.DIJKSTRA_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 
-public class ShortestPathDijkstraMutateProc extends BaseProc {
+public class ShortestPathDijkstraMutateProc {
     @Context
     public GraphDataScience facade;
 
@@ -54,10 +53,9 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
         @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
     ) {
-        return new MemoryEstimationExecutor<>(
-            new ShortestPathDijkstraMutateSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphNameOrConfiguration, algoConfiguration);
+        return facade.pathFinding().singlePairShortestPathDijkstraMutateEstimate(
+            graphNameOrConfiguration,
+            algoConfiguration
+        );
     }
 }
diff --git a/...inding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathDijkstraStreamProc.java b/...inding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathDijkstraStreamProc.java
@@ -19,10 +19,8 @@
  */
 package org.neo4j.gds.paths.sourcetarget;
 
-import org.neo4j.gds.BaseProc;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
-import org.neo4j.gds.procedures.pathfinding.PathFindingStreamResult;
 import org.neo4j.gds.procedures.GraphDataScience;
+import org.neo4j.gds.procedures.pathfinding.PathFindingStreamResult;
 import org.neo4j.gds.results.MemoryEstimateResult;
 import org.neo4j.procedure.Context;
 import org.neo4j.procedure.Description;
@@ -32,10 +30,11 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.neo4j.gds.ProcedureConstants.ESTIMATE_DESCRIPTION;
 import static org.neo4j.gds.paths.sourcetarget.SinglePairShortestPathConstants.DIJKSTRA_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 
-public class ShortestPathDijkstraStreamProc extends BaseProc {
+public class ShortestPathDijkstraStreamProc {
     @Context
     public GraphDataScience facade;
 
@@ -54,13 +53,9 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
         @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
     ) {
-
-        return new MemoryEstimationExecutor<>(
-            new ShortestPathDijkstraStreamSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphNameOrConfiguration, algoConfiguration);
+        return facade.pathFinding().singlePairShortestPathDijkstraStreamEstimate(
+            graphNameOrConfiguration,
+            algoConfiguration
+        );
     }
-
-
 }
diff --git a/...finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathDijkstraWriteProc.java b/...finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathDijkstraWriteProc.java
@@ -19,10 +19,6 @@
  */
 package org.neo4j.gds.paths.sourcetarget;
 
-import org.neo4j.gds.BaseProc;
-import org.neo4j.gds.core.write.RelationshipStreamExporterBuilder;
-import org.neo4j.gds.executor.ExecutionContext;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.procedures.GraphDataScience;
 import org.neo4j.gds.results.MemoryEstimateResult;
 import org.neo4j.gds.results.StandardWriteRelationshipsResult;
@@ -34,17 +30,15 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.neo4j.gds.ProcedureConstants.ESTIMATE_DESCRIPTION;
 import static org.neo4j.gds.paths.sourcetarget.SinglePairShortestPathConstants.DIJKSTRA_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 import static org.neo4j.procedure.Mode.WRITE;
 
-public class ShortestPathDijkstraWriteProc extends BaseProc {
+public class ShortestPathDijkstraWriteProc {
     @Context
     public GraphDataScience facade;
 
-    @Context
-    public RelationshipStreamExporterBuilder relationshipStreamExporterBuilder;
-
     @Procedure(name = "gds.shortestPath.dijkstra.write", mode = WRITE)
     @Description(DIJKSTRA_DESCRIPTION)
     public Stream<StandardWriteRelationshipsResult> write(
@@ -60,16 +54,9 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
         @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
     ) {
-        return new MemoryEstimationExecutor<>(
-            new ShortestPathDijkstraWriteSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphNameOrConfiguration, algoConfiguration);
+        return facade.pathFinding().singlePairShortestPathDijkstraWriteEstimate(
+            graphNameOrConfiguration,
+            algoConfiguration
+        );
     }
-
-    @Override
-    public ExecutionContext executionContext() {
-        return super.executionContext().withRelationshipStreamExporterBuilder(relationshipStreamExporterBuilder);
-    }
-
 }
diff --git a/...th-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathYensMutateProc.java b/...th-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathYensMutateProc.java
@@ -19,8 +19,6 @@
  */
 package org.neo4j.gds.paths.sourcetarget;
 
-import org.neo4j.gds.BaseProc;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.procedures.GraphDataScience;
 import org.neo4j.gds.procedures.pathfinding.PathFindingMutateResult;
 import org.neo4j.gds.results.MemoryEstimateResult;
@@ -32,10 +30,11 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.neo4j.gds.ProcedureConstants.ESTIMATE_DESCRIPTION;
 import static org.neo4j.gds.paths.sourcetarget.SinglePairShortestPathConstants.YENS_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 
-public class ShortestPathYensMutateProc extends BaseProc {
+public class ShortestPathYensMutateProc {
     @Context
     public GraphDataScience facade;
 
@@ -54,10 +53,6 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphNameOrConfiguration") Object graphName,
         @Name(value = "algoConfiguration") Map<String, Object> configuration
     ) {
-        return new MemoryEstimationExecutor<>(
-            new ShortestPathYensMutateSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphName, configuration);
+        return facade.pathFinding().singlePairShortestPathYensMutateEstimate(graphName, configuration);
     }
 }
diff --git a/...th-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathYensStreamProc.java b/...th-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathYensStreamProc.java
@@ -19,8 +19,6 @@
  */
 package org.neo4j.gds.paths.sourcetarget;
 
-import org.neo4j.gds.BaseProc;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.procedures.GraphDataScience;
 import org.neo4j.gds.procedures.pathfinding.PathFindingStreamResult;
 import org.neo4j.gds.results.MemoryEstimateResult;
@@ -32,10 +30,11 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.neo4j.gds.ProcedureConstants.ESTIMATE_DESCRIPTION;
 import static org.neo4j.gds.paths.sourcetarget.SinglePairShortestPathConstants.YENS_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 
-public class ShortestPathYensStreamProc extends BaseProc {
+public class ShortestPathYensStreamProc {
     @Context
     public GraphDataScience facade;
 
@@ -54,10 +53,6 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphName") Object graphName,
         @Name(value = "configuration") Map<String, Object> configuration
     ) {
-        return new MemoryEstimationExecutor<>(
-            new ShortestPathYensStreamSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphName, configuration);
+        return facade.pathFinding().singlePairShortestPathYensStreamEstimate(graphName, configuration);
     }
 }
diff --git a/...ath-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathYensWriteProc.java b/...ath-finding/src/main/java/org/neo4j/gds/paths/sourcetarget/ShortestPathYensWriteProc.java
@@ -19,10 +19,6 @@
  */
 package org.neo4j.gds.paths.sourcetarget;
 
-import org.neo4j.gds.BaseProc;
-import org.neo4j.gds.core.write.RelationshipStreamExporterBuilder;
-import org.neo4j.gds.executor.ExecutionContext;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.procedures.GraphDataScience;
 import org.neo4j.gds.results.MemoryEstimateResult;
 import org.neo4j.gds.results.StandardWriteRelationshipsResult;
@@ -34,17 +30,15 @@
 import java.util.Map;
 import java.util.stream.Stream;
 
+import static org.neo4j.gds.ProcedureConstants.ESTIMATE_DESCRIPTION;
 import static org.neo4j.gds.paths.sourcetarget.SinglePairShortestPathConstants.YENS_DESCRIPTION;
 import static org.neo4j.procedure.Mode.READ;
 import static org.neo4j.procedure.Mode.WRITE;
 
-public class ShortestPathYensWriteProc extends BaseProc {
+public class ShortestPathYensWriteProc {
     @Context
     public GraphDataScience facade;
 
-    @Context
-    public RelationshipStreamExporterBuilder relationshipStreamExporterBuilder;
-
     @Procedure(name = "gds.shortestPath.yens.write", mode = WRITE)
     @Description(YENS_DESCRIPTION)
     public Stream<StandardWriteRelationshipsResult> write(
@@ -60,15 +54,6 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphName") Object graphName,
         @Name(value = "configuration") Map<String, Object> configuration
     ) {
-        return new MemoryEstimationExecutor<>(
-            new ShortestPathYensWriteSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphName, configuration);
-    }
-
-    @Override
-    public ExecutionContext executionContext() {
-        return super.executionContext().withRelationshipStreamExporterBuilder(relationshipStreamExporterBuilder);
+        return facade.pathFinding().singlePairShortestPathYensWriteEstimate(graphName, configuration);
     }
 }