Fix graphsage memory estimation

Author: IoannisPanagiotas

algo/src/main/java/org/neo4j/gds/algorithms/embeddings/NodeEmbeddingsAlgorithmsEstimateBusinessFacade.java

@@ -30,18 +30,23 @@
 import org.neo4j.gds.embeddings.hashgnn.HashGNNFactory;
 import org.neo4j.gds.embeddings.node2vec.Node2VecAlgorithmFactory;
 import org.neo4j.gds.embeddings.node2vec.Node2VecBaseConfig;
+import org.neo4j.gds.modelcatalogservices.ModelCatalogService;
 import org.neo4j.gds.results.MemoryEstimateResult;
 
 import java.util.Optional;
 
 public class NodeEmbeddingsAlgorithmsEstimateBusinessFacade {
 
     private final AlgorithmEstimator algorithmEstimator;
+    private final ModelCatalogService modelCatalogService;
+
 
     public NodeEmbeddingsAlgorithmsEstimateBusinessFacade(
-        AlgorithmEstimator algorithmEstimator
+        AlgorithmEstimator algorithmEstimator,
+        ModelCatalogService modelCatalogService
     ) {
         this.algorithmEstimator = algorithmEstimator;
+        this.modelCatalogService = modelCatalogService;
     }
 
     public <C extends Node2VecBaseConfig> MemoryEstimateResult node2Vec(
@@ -64,7 +69,7 @@ public <C extends GraphSageBaseConfig> MemoryEstimateResult graphSage(
             graphNameOrConfiguration,
             configuration,
             Optional.empty(),
-            new GraphSageAlgorithmFactory(null)
+            new GraphSageAlgorithmFactory(modelCatalogService.get())
         );
     }
 

proc/embeddings/src/test/java/org/neo4j/gds/embeddings/graphsage/GraphSageMutateProcTest.java

@@ -48,6 +48,7 @@
 import java.util.stream.Stream;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatNoException;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.neo4j.gds.ElementProjection.PROJECT_ALL;
@@ -262,4 +263,32 @@ private static Stream<Arguments> missingNodeProperties() {
         );
     }
 
+    @Test
+    void shouldEstimateMemory() {
+        var trainQuery = GdsCypher.call(graphName)
+            .algo("gds.beta.graphSage")
+            .trainMode()
+            .addParameter("sampleSizes", List.of(2, 4))
+            .addParameter("featureProperties", List.of("age", "birth_year", "death_year"))
+            .addParameter("embeddingDimension", 16)
+            .addParameter("activationFunction", ActivationFunction.SIGMOID)
+            .addParameter("aggregator", "mean")
+            .addParameter("modelName", modelName)
+            .yields();
+
+        runQuery(trainQuery);
+
+        var mutatePropertyKey = "embedding";
+        var query = GdsCypher.call("embeddingsGraph")
+            .algo("gds.beta.graphSage")
+            .mutateEstimation()
+            .addParameter("mutateProperty", mutatePropertyKey)
+            .addParameter("modelName", modelName)
+            .yields("requiredMemory");
+
+        assertThatNoException().isThrownBy(() -> runQuery(query));
+
+
+    }
+
 }

proc/embeddings/src/test/java/org/neo4j/gds/embeddings/graphsage/GraphSageStreamProcTest.java

@@ -177,6 +177,7 @@ void weightedGraphSage() {
             .addParameter("concurrency", 1)
             .addParameter("modelName", modelName)
             .yields();
+
         long[] nodeIds = idFunction.of("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o");
         var rowCount = runQueryWithRowConsumer(streamQuery, (resultRow) -> {
             assertThat(nodeIds).contains((long) resultRow.getNumber("nodeId"));

procedures/integration/src/main/java/org/neo4j/gds/procedures/integration/AlgorithmProcedureFacadeProvider.java

@@ -294,7 +294,10 @@ NodeEmbeddingsProcedureFacade createNodeEmbeddingsProcedureFacade() {
             writeNodePropertyService
         );
 
-        var estimateBusinessFacade = new NodeEmbeddingsAlgorithmsEstimateBusinessFacade(algorithmEstimator);
+        var estimateBusinessFacade = new NodeEmbeddingsAlgorithmsEstimateBusinessFacade(
+            algorithmEstimator,
+            modelCatalogService
+        );
 
         // procedure facade
         return new NodeEmbeddingsProcedureFacade(