migrate node classification predict estimate

Author: lassewesth

proc/community/src/integrationTest/java/org/neo4j/gds/labelpropagation/LabelPropagationMutateProcTest.java

@@ -73,6 +73,7 @@
 import org.neo4j.gds.metrics.procedures.DeprecatedProceduresMetricService;
 import org.neo4j.gds.procedures.GraphDataScienceProcedures;
 import org.neo4j.gds.procedures.GraphDataScienceProceduresBuilder;
+import org.neo4j.gds.procedures.algorithms.AlgorithmsProcedureFacade;
 import org.neo4j.gds.procedures.algorithms.community.CommunityProcedureFacade;
 import org.neo4j.gds.procedures.algorithms.configuration.ConfigurationParser;
 import org.neo4j.gds.procedures.algorithms.configuration.UserSpecificConfigurationParser;
@@ -502,11 +503,7 @@ private GraphDataScienceProcedures constructFacade() {
 
         var configurationParser = new UserSpecificConfigurationParser(new ConfigurationParser(DefaultsConfiguration.Instance, LimitsConfiguration.Instance),requestScopedDependencies.getUser());
 
-        var genericStub = GenericStub.create(
-            graphStoreCatalogService,
-            configurationParser,
-            requestScopedDependencies
-        );
+        var genericStub = new GenericStub(configurationParser, null);
         var applicationsFacade = ApplicationsFacade.create(
             logMock,
             null,
@@ -537,7 +534,7 @@ private GraphDataScienceProcedures constructFacade() {
         );
 
         return new GraphDataScienceProceduresBuilder(Log.noOpLog())
-            .with(communityProcedureFacade)
+            .with(new AlgorithmsProcedureFacade(null, communityProcedureFacade, null, null, null, null, null))
             .with(DeprecatedProceduresMetricService.PASSTHROUGH)
             .build();
     }

proc/community/src/integrationTest/java/org/neo4j/gds/modularityoptimization/ModularityOptimizationMutateProcTest.java

@@ -73,6 +73,7 @@
 import org.neo4j.gds.metrics.projections.ProjectionMetricsService;
 import org.neo4j.gds.procedures.GraphDataScienceProcedures;
 import org.neo4j.gds.procedures.GraphDataScienceProceduresBuilder;
+import org.neo4j.gds.procedures.algorithms.AlgorithmsProcedureFacade;
 import org.neo4j.gds.procedures.algorithms.community.CommunityProcedureFacade;
 import org.neo4j.gds.procedures.algorithms.configuration.ConfigurationParser;
 import org.neo4j.gds.procedures.algorithms.configuration.UserSpecificConfigurationParser;
@@ -509,11 +510,7 @@ private GraphDataScienceProcedures createFacade() {
 
         var configurationParser = new UserSpecificConfigurationParser(new ConfigurationParser(DefaultsConfiguration.Instance, LimitsConfiguration.Instance),requestScopedDependencies.getUser());
 
-        var genericStub = GenericStub.create(
-            graphStoreCatalogService,
-            configurationParser,
-            requestScopedDependencies
-        );
+        var genericStub = new GenericStub(configurationParser, null);
         var applicationsFacade = ApplicationsFacade.create(
             logMock,
             null,
@@ -544,7 +541,7 @@ private GraphDataScienceProcedures createFacade() {
         );
 
         return new GraphDataScienceProceduresBuilder(logMock)
-            .with(communityProcedureFacade)
+            .with(new AlgorithmsProcedureFacade(null, communityProcedureFacade, null, null, null, null, null))
             .with(DeprecatedProceduresMetricService.PASSTHROUGH)
             .build();
     }

proc/community/src/integrationTest/java/org/neo4j/gds/wcc/WccMutateProcTest.java

@@ -74,6 +74,7 @@
 import org.neo4j.gds.metrics.projections.ProjectionMetricsService;
 import org.neo4j.gds.procedures.GraphDataScienceProcedures;
 import org.neo4j.gds.procedures.GraphDataScienceProceduresBuilder;
+import org.neo4j.gds.procedures.algorithms.AlgorithmsProcedureFacade;
 import org.neo4j.gds.procedures.algorithms.community.CommunityProcedureFacade;
 import org.neo4j.gds.procedures.algorithms.community.WccMutateResult;
 import org.neo4j.gds.procedures.algorithms.configuration.ConfigurationParser;
@@ -605,11 +606,7 @@ private GraphDataScienceProcedures constructGraphDataScienceProcedures() {
 
         var configurationParser = new UserSpecificConfigurationParser(new ConfigurationParser(DefaultsConfiguration.Instance, LimitsConfiguration.Instance),requestScopedDependencies.getUser());
 
-        var genericStub = GenericStub.create(
-            graphStoreCatalogService,
-            configurationParser,
-            requestScopedDependencies
-        );
+        var genericStub = new GenericStub(configurationParser, null);
         var communityProcedureFacade = CommunityProcedureFacade.create(
             genericStub,
             applicationsFacade,
@@ -622,7 +619,7 @@ private GraphDataScienceProcedures constructGraphDataScienceProcedures() {
         );
 
         return new GraphDataScienceProceduresBuilder(Log.noOpLog())
-            .with(communityProcedureFacade)
+            .with(new AlgorithmsProcedureFacade(null, communityProcedureFacade, null, null, null, null, null))
             .with(DeprecatedProceduresMetricService.PASSTHROUGH)
             .build();
     }

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/NodePropertyPredictPipelineFilterUtil.java

@@ -24,6 +24,7 @@
 import org.neo4j.gds.core.model.ModelCatalog;
 import org.neo4j.gds.ml.models.BaseModelData;
 import org.neo4j.gds.ml.pipeline.nodePipeline.NodePropertyPipelineBaseTrainConfig;
+import org.neo4j.gds.procedures.pipelines.NodePropertyPredictPipelineBaseConfig;
 
 import java.util.List;
 

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPipelineMutateProc.java

@@ -20,12 +20,12 @@
 package org.neo4j.gds.ml.pipeline.node.classification.predict;
 
 import org.neo4j.gds.BaseProc;
+import org.neo4j.gds.applications.algorithms.machinery.MemoryEstimateResult;
 import org.neo4j.gds.core.model.ModelCatalog;
 import org.neo4j.gds.executor.ExecutionContext;
-import org.neo4j.gds.executor.MemoryEstimationExecutor;
 import org.neo4j.gds.executor.ProcedureExecutor;
 import org.neo4j.gds.ml.pipeline.node.PredictMutateResult;
-import org.neo4j.gds.applications.algorithms.machinery.MemoryEstimateResult;
+import org.neo4j.gds.procedures.GraphDataScienceProcedures;
 import org.neo4j.procedure.Context;
 import org.neo4j.procedure.Description;
 import org.neo4j.procedure.Mode;
@@ -43,6 +43,9 @@ public class NodeClassificationPipelineMutateProc extends BaseProc {
     @Context
     public ModelCatalog internalModelCatalog;
 
+    @Context
+    public GraphDataScienceProcedures facade;
+
     @Procedure(name = "gds.beta.pipeline.nodeClassification.predict.mutate", mode = Mode.READ)
     @Description(PREDICT_DESCRIPTION)
     public Stream<PredictMutateResult> mutate(
@@ -59,15 +62,10 @@ public Stream<PredictMutateResult> mutate(
     @Procedure(name = "gds.beta.pipeline.nodeClassification.predict.mutate.estimate", mode = Mode.READ)
     @Description(ESTIMATE_PREDICT_DESCRIPTION)
     public Stream<MemoryEstimateResult> estimate(
-        @Name(value = "graphName") Object graphName,
-        @Name(value = "configuration") Map<String, Object> configuration
+        @Name(value = "graphName") Object graphNameOrConfiguration,
+        @Name(value = "configuration") Map<String, Object> algorithmConfiguration
     ) {
-        preparePipelineConfig(graphName, configuration);
-        return new MemoryEstimationExecutor<>(
-            new NodeClassificationPipelineMutateSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphName, configuration);
+        return facade.pipelines().nodeClassificationMutateEstimate(graphNameOrConfiguration, algorithmConfiguration);
     }
 
     @Override

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPipelineMutateSpec.java

@@ -28,6 +28,8 @@
 import org.neo4j.gds.executor.GdsCallable;
 import org.neo4j.gds.procedures.algorithms.configuration.NewConfigFunction;
 import org.neo4j.gds.ml.pipeline.node.PredictMutateResult;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictConfigPreProcessor;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictPipelineMutateConfig;
 import org.neo4j.gds.result.AbstractResultBuilder;
 
 import java.util.Map;

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPipelineWriteSpec.java

@@ -28,6 +28,7 @@
 import org.neo4j.gds.executor.ExecutionContext;
 import org.neo4j.gds.executor.GdsCallable;
 import org.neo4j.gds.procedures.algorithms.configuration.NewConfigFunction;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictConfigPreProcessor;
 
 import java.util.List;
 import java.util.Map;

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictNewMutateConfigFn.java

@@ -22,6 +22,8 @@
 import org.neo4j.gds.core.CypherMapWrapper;
 import org.neo4j.gds.core.model.ModelCatalog;
 import org.neo4j.gds.procedures.algorithms.configuration.NewConfigFunction;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictPipelineMutateConfig;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictPipelineMutateConfigImpl;
 
 import static org.neo4j.gds.ml.pipeline.node.NodePropertyPredictPipelineFilterUtil.generatePredictPipelineFilter;
 

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineAlgorithmFactory.java

@@ -32,6 +32,7 @@
 import org.neo4j.gds.ml.models.Classifier;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.train.NodeClassificationPipelineModelInfo;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.train.NodeClassificationPipelineTrainConfig;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictPipelineBaseConfig;
 
 public class NodeClassificationPredictPipelineAlgorithmFactory
     <CONFIG extends NodeClassificationPredictPipelineBaseConfig>

proc/machine-learning/src/main/java/org/neo4j/gds/ml/pipeline/node/classification/predict/NodeClassificationPredictPipelineExecutor.java

@@ -24,7 +24,6 @@
 import org.neo4j.gds.core.model.Model;
 import org.neo4j.gds.core.model.ModelCatalog;
 import org.neo4j.gds.mem.MemoryEstimation;
-import org.neo4j.gds.mem.MemoryEstimations;
 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;
@@ -42,18 +41,18 @@
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.train.NodeClassificationPipelineModelInfo;
 import org.neo4j.gds.ml.pipeline.nodePipeline.classification.train.NodeClassificationPipelineTrainConfig;
 import org.neo4j.gds.procedures.algorithms.AlgorithmsProcedureFacade;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictPipelineBaseConfig;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictPipelineConstants;
+import org.neo4j.gds.procedures.pipelines.NodeClassificationPredictPipelineEstimator;
 import org.neo4j.gds.utils.StringJoining;
 
-import java.util.List;
-
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
 
 public class NodeClassificationPredictPipelineExecutor extends PredictPipelineExecutor<
     NodeClassificationPredictPipelineBaseConfig,
     NodePropertyPredictPipeline,
     NodeClassificationPipelineResult
     > {
-    private static final int MIN_BATCH_SIZE = 100;
     private final Classifier.ClassifierData modelData;
     private final LocalIdMap classIdMap;
     private final PipelineGraphFilter predictGraphFilter;
@@ -90,35 +89,9 @@ public static MemoryEstimation estimate(
         ModelCatalog modelCatalog,
         AlgorithmsProcedureFacade algorithmsProcedureFacade
     ) {
-        var pipeline = model.customInfo().pipeline();
-        var classCount = model.customInfo().classes().size();
-        var featureCount = model.data().featureDimension();
-
-        var combinedNodeLabels = configuration.targetNodeLabels().isEmpty() ? model.trainConfig().targetNodeLabels() : configuration.targetNodeLabels();
-        var combinedRelationshipTypes = configuration.relationshipTypes().isEmpty() ? model.trainConfig().relationshipTypes() : configuration.relationshipTypes();
-
-        MemoryEstimation nodePropertyStepEstimation = NodePropertyStepExecutor.estimateNodePropertySteps(
-            algorithmsProcedureFacade,
-            modelCatalog,
-            configuration.username(),
-            pipeline.nodePropertySteps(),
-            combinedNodeLabels,
-            combinedRelationshipTypes
-        );
-
-        var predictionEstimation = MemoryEstimations.builder().add(
-            "Pipeline Predict",
-            NodeClassificationPredict.memoryEstimationWithDerivedBatchSize(
-                model.data().trainerMethod(),
-                configuration.includePredictedProbabilities(),
-                MIN_BATCH_SIZE,
-                featureCount,
-                classCount,
-                false
-            )
-        ).build();
+        var estimator = new NodeClassificationPredictPipelineEstimator(modelCatalog, algorithmsProcedureFacade);
 
-        return MemoryEstimations.maxEstimation(List.of(nodePropertyStepEstimation, predictionEstimation));
+        return estimator.estimate(model, configuration);
     }
 
     @Override
@@ -143,7 +116,7 @@ protected NodeClassificationPipelineResult execute() {
         var nodeClassificationResult =  new NodeClassificationPredict(
             ClassifierFactory.create(modelData),
             features,
-            MIN_BATCH_SIZE,
+            NodeClassificationPredictPipelineConstants.MIN_BATCH_SIZE,
             config.concurrency(),
             config.includePredictedProbabilities(),
             progressTracker,