Merge pull request #9751 from lassewesth/lpppp4

migrate link prediction stream estimate

Author: lassewesth

proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/LinkPredictionPipelineStreamProc.java

@@ -19,12 +19,9 @@
  */
 package org.neo4j.gds.ml.linkmodels.pipeline.predict;
 
-import org.neo4j.gds.BaseProc;
-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.applications.algorithms.machinery.MemoryEstimateResult;
+import org.neo4j.gds.procedures.GraphDataScienceProcedures;
+import org.neo4j.gds.procedures.pipelines.StreamResult;
 import org.neo4j.procedure.Context;
 import org.neo4j.procedure.Description;
 import org.neo4j.procedure.Mode;
@@ -36,24 +33,18 @@
 
 import static org.neo4j.gds.ml.linkmodels.pipeline.LinkPredictionPipelineCompanion.ESTIMATE_PREDICT_DESCRIPTION;
 import static org.neo4j.gds.ml.linkmodels.pipeline.LinkPredictionPipelineCompanion.PREDICT_DESCRIPTION;
-import static org.neo4j.gds.ml.pipeline.PipelineCompanion.preparePipelineConfig;
-
-public class LinkPredictionPipelineStreamProc extends BaseProc {
 
+public class LinkPredictionPipelineStreamProc {
     @Context
-    public ModelCatalog modelCatalog;
+    public GraphDataScienceProcedures facade;
 
     @Procedure(name = "gds.beta.pipeline.linkPrediction.predict.stream", mode = Mode.READ)
     @Description(PREDICT_DESCRIPTION)
     public Stream<StreamResult> stream(
         @Name(value = "graphName") String graphName,
         @Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration
     ) {
-        preparePipelineConfig(graphName, configuration);
-        return new ProcedureExecutor<>(
-            new LinkPredictionPipelineStreamSpec(),
-            executionContext()
-        ).compute(graphName, configuration);
+        return facade.pipelines().linkPrediction().stream(graphName, configuration);
     }
 
     @Procedure(name = "gds.beta.pipeline.linkPrediction.predict.stream.estimate", mode = Mode.READ)
@@ -62,17 +53,6 @@ public Stream<MemoryEstimateResult> estimate(
         @Name(value = "graphNameOrConfiguration") Object graphNameOrConfiguration,
         @Name(value = "algoConfiguration") Map<String, Object> algoConfiguration
     ) {
-        preparePipelineConfig(graphNameOrConfiguration, algoConfiguration);
-        return new MemoryEstimationExecutor<>(
-            new LinkPredictionPipelineStreamSpec(),
-            executionContext(),
-            transactionContext()
-        ).computeEstimate(graphNameOrConfiguration, algoConfiguration);
-    }
-
-    @Override
-    public ExecutionContext executionContext() {
-        return super.executionContext().withModelCatalog(modelCatalog);
+        return facade.pipelines().linkPrediction().streamEstimate(graphNameOrConfiguration, algoConfiguration);
     }
-
 }

proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/LinkPredictionPipelineStreamSpec.java

@@ -28,6 +28,7 @@
 import org.neo4j.gds.ml.linkmodels.LinkPredictionResult;
 import org.neo4j.gds.procedures.pipelines.LinkPredictionPredictPipelineExecutor;
 import org.neo4j.gds.procedures.pipelines.LinkPredictionPredictPipelineStreamConfig;
+import org.neo4j.gds.procedures.pipelines.StreamResult;
 
 import java.util.Collection;
 import java.util.stream.Stream;

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/LinkPredictionFacade.java

@@ -169,4 +169,31 @@ public Stream<MemoryEstimateResult> mutateEstimate(
 
         return Stream.of(result);
     }
+
+    public Stream<StreamResult> stream(
+        String graphNameAsString,
+        Map<String, Object> configuration
+    ) {
+        PipelineCompanion.preparePipelineConfig(graphNameAsString, configuration);
+
+        var graphName = GraphName.parse(graphNameAsString);
+
+        return pipelineApplications.linkPredictionStream(graphName, configuration);
+    }
+
+    public Stream<MemoryEstimateResult> streamEstimate(
+        Object graphNameOrConfiguration,
+        Map<String, Object> rawConfiguration
+    ) {
+        PipelineCompanion.preparePipelineConfig(graphNameOrConfiguration, rawConfiguration);
+
+        var configuration = pipelineConfigurationParser.parseLinkPredictionPredictPipelineStreamConfig(rawConfiguration);
+
+        var result = pipelineApplications.linkPredictionEstimate(
+            graphNameOrConfiguration,
+            configuration
+        );
+
+        return Stream.of(result);
+    }
 }

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/LinkPredictionPipelineStreamResultBuilder.java

@@ -0,0 +1,75 @@
+/*
+ * 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.procedures.pipelines;
+
+import org.neo4j.gds.api.Graph;
+import org.neo4j.gds.api.GraphStore;
+import org.neo4j.gds.applications.algorithms.machinery.StreamResultBuilder;
+import org.neo4j.gds.logging.Log;
+import org.neo4j.gds.ml.linkmodels.LinkPredictionResult;
+
+import java.util.Optional;
+import java.util.stream.Stream;
+
+class LinkPredictionPipelineStreamResultBuilder implements StreamResultBuilder<LinkPredictionResult, StreamResult> {
+    private final Log log;
+    private final TrainedLPPipelineModel trainedLPPipelineModel;
+    private final LinkPredictionPredictPipelineStreamConfig configuration;
+
+    public LinkPredictionPipelineStreamResultBuilder(
+        Log log,
+        TrainedLPPipelineModel trainedLPPipelineModel,
+        LinkPredictionPredictPipelineStreamConfig configuration
+    ) {
+        this.log = log;
+        this.trainedLPPipelineModel = trainedLPPipelineModel;
+        this.configuration = configuration;
+    }
+
+    @Override
+    public Stream<StreamResult> build(
+        Graph graph,
+        GraphStore graphStore,
+        Optional<LinkPredictionResult> result
+    ) {
+        if (result.isEmpty()) return Stream.empty();
+
+        var linkPredictionResult = result.get();
+
+        var model = trainedLPPipelineModel.get(
+            configuration.modelName(),
+            configuration.username()
+        );
+
+        var lpGraphStoreFilter = LPGraphStoreFilterFactory.generate(
+            log, model.trainConfig(),
+            configuration,
+            graphStore
+        );
+
+        var filteredGraph = graphStore.getGraph(lpGraphStoreFilter.predictNodeLabels());
+
+        return linkPredictionResult.stream().map(predictedLink -> new StreamResult(
+            filteredGraph.toOriginalNodeId(predictedLink.sourceId()),
+            filteredGraph.toOriginalNodeId(predictedLink.targetId()),
+            predictedLink.probability()
+        ));
+    }
+}

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/PipelineApplications.java

@@ -389,6 +389,24 @@ MutateResult linkPredictionMutate(GraphName graphName, Map<String, Object> rawCo
         );
     }
 
+    Stream<StreamResult> linkPredictionStream(GraphName graphName, Map<String, Object> rawConfiguration) {
+        var configuration = pipelineConfigurationParser.parseLinkPredictionPredictPipelineStreamConfig(rawConfiguration);
+        var label = new StandardLabel("LinkPredictionPipelineStream");
+        var computation = constructLinkPredictionComputation(configuration, label);
+        var resultBuilder = new LinkPredictionPipelineStreamResultBuilder(log, trainedLPPipelineModel, configuration);
+
+        return algorithmProcessingTemplate.processAlgorithmForStream(
+            Optional.empty(),
+            graphName,
+            configuration,
+            Optional.empty(),
+            label,
+            () -> linkPredictionMemoryEstimation(configuration),
+            computation,
+            resultBuilder
+        );
+    }
+
     MemoryEstimateResult nodeClassificationPredictEstimate(
         Object graphNameOrConfiguration,
         NodeClassificationPredictPipelineBaseConfig configuration

procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/PipelineConfigurationParser.java

@@ -72,6 +72,10 @@ LinkPredictionSplitConfig parseLinkPredictionSplitConfig(Map<String, Object> raw
         );
     }
 
+    LinkPredictionPredictPipelineStreamConfig parseLinkPredictionPredictPipelineStreamConfig(Map<String, Object> configuration) {
+        return parseConfiguration(LinkPredictionPredictPipelineStreamConfig::of, configuration);
+    }
+
     TunableTrainerConfig parseLogisticRegressionTrainerConfig(Map<String, Object> configuration) {
         return parseTrainerConfiguration(
             configuration,

proc/machine-learning/src/main/java/org/neo4j/gds/ml/linkmodels/pipeline/predict/StreamResult.java renamed to procedures/pipelines-facade/src/main/java/org/neo4j/gds/procedures/pipelines/StreamResult.java

@@ -17,11 +17,9 @@
  * 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.ml.linkmodels.pipeline.predict;
+package org.neo4j.gds.procedures.pipelines;
 
-@SuppressWarnings("unused")
 public final class StreamResult {
-
     public final long node1;
     public final long node2;
     public final double probability;