Merge pull request #10730 from IoannisPanagiotas/write-service-more

Write service more

Author: IoannisPanagiotas

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/Neo4jDatabaseRelationshipWriter.java

@@ -19,24 +19,32 @@
  */
 package org.neo4j.gds.applications.algorithms.machinery;
 
+import org.neo4j.gds.api.ExportedRelationship;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.api.ResultStore;
+import org.neo4j.gds.api.nodeproperties.ValueType;
 import org.neo4j.gds.api.properties.relationships.RelationshipWithPropertyConsumer;
 import org.neo4j.gds.applications.algorithms.metadata.RelationshipsWritten;
+import org.neo4j.gds.core.concurrency.Concurrency;
 import org.neo4j.gds.core.utils.logging.LoggerForProgressTrackingAdapter;
 import org.neo4j.gds.core.utils.progress.JobId;
 import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.tasks.TaskProgressTracker;
 import org.neo4j.gds.core.write.RelationshipExporter;
 import org.neo4j.gds.core.write.RelationshipExporterBuilder;
+import org.neo4j.gds.core.write.RelationshipStreamExporter;
+import org.neo4j.gds.core.write.RelationshipStreamExporterBuilder;
 import org.neo4j.gds.logging.Log;
 import org.neo4j.gds.termination.TerminationFlag;
 
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 final class Neo4jDatabaseRelationshipWriter {
-    static RelationshipsWritten writeRelationship(
+
+    static RelationshipsWritten writeRelationshipsFromGraph(
         String writeRelationshipType,
         String writeProperty,
         TaskRegistryFactory taskRegistryFactory,
@@ -49,24 +57,23 @@ static RelationshipsWritten writeRelationship(
         Optional<ResultStore> resultStore,
         RelationshipWithPropertyConsumer relationshipConsumer,
         JobId jobId
-    ){
+    ) {
 
         var progressTracker = new TaskProgressTracker(
-                RelationshipExporter.baseTask(taskName, graph.relationshipCount()),
-                new LoggerForProgressTrackingAdapter(log),
-                RelationshipExporterBuilder.TYPED_DEFAULT_WRITE_CONCURRENCY,
-                taskRegistryFactory
-            );
-
-            var exporter = relationshipExporterBuilder
-                .withIdMappingOperator(rootIdMap::toOriginalNodeId)
-                .withGraph(graph)
-                .withTerminationFlag(algorithmTerminationFlag)
-                .withProgressTracker(progressTracker)
-                .withResultStore(resultStore)
-                .withJobId(jobId)
-                .build();
+            RelationshipExporter.baseTask(taskName, graph.relationshipCount()),
+            new LoggerForProgressTrackingAdapter(log),
+            RelationshipExporterBuilder.TYPED_DEFAULT_WRITE_CONCURRENCY,
+            taskRegistryFactory
+        );
 
+        var exporter = relationshipExporterBuilder
+            .withIdMappingOperator(rootIdMap::toOriginalNodeId)
+            .withGraph(graph)
+            .withTerminationFlag(algorithmTerminationFlag)
+            .withProgressTracker(progressTracker)
+            .withResultStore(resultStore)
+            .withJobId(jobId)
+            .build();
 
         try {
             exporter.write(
@@ -82,8 +89,58 @@ static RelationshipsWritten writeRelationship(
         }
 
         return new RelationshipsWritten(graph.relationshipCount());
+    }
 
+    static RelationshipsWritten writeRelationshipsFromStream(
+        String writeRelationshipType,
+        List<String> properties,
+        List<ValueType> valueTypes,
+        TaskRegistryFactory taskRegistryFactory,
+        RelationshipStreamExporterBuilder relationshipStreamExporterBuilder,
+        Stream<ExportedRelationship> relationshipStream,
+        IdMap rootIdMap,
+        Log log,
+        String taskName,
+        TerminationFlag algorithmTerminationFlag,
+        Optional<ResultStore> maybeResultStore,
+        JobId jobId
+    ) {
+
+        var progressTracker = new TaskProgressTracker(
+            RelationshipStreamExporter.baseTask(taskName),
+            new LoggerForProgressTrackingAdapter(log),
+            new Concurrency(1),
+            taskRegistryFactory
+        );
+
+        // When we are writing to the result store, the result stream might not be consumed
+        // inside the current transaction. This causes the stream to immediately return an empty stream
+        // as the termination flag, which is bound to the current transaction is set to true. We therefore
+        // need to collect the stream and trigger an eager computation.
+        var maybeCollectedStream = maybeResultStore
+            .map(__ -> relationshipStream.toList().stream())
+            .orElse(relationshipStream);
+
+        // configure the exporter
+        var relationshipStreamExporter = relationshipStreamExporterBuilder
+            .withResultStore(maybeResultStore)
+            .withIdMappingOperator(rootIdMap::toOriginalNodeId)
+            .withProgressTracker(progressTracker)
+            .withRelationships(maybeCollectedStream)
+            .withTerminationFlag(algorithmTerminationFlag)
+            .withJobId(jobId)
+            .build();
+
+        var relationshipsWritten = relationshipStreamExporter.write(
+            writeRelationshipType,
+            properties,
+            valueTypes
+        );
+
+        return new RelationshipsWritten(relationshipsWritten);
     }
 
+
     private Neo4jDatabaseRelationshipWriter() {}
 }
+

applications/algorithms/machinery/src/main/java/org/neo4j/gds/applications/algorithms/machinery/WriteRelationshipService.java

@@ -19,15 +19,19 @@
  */
 package org.neo4j.gds.applications.algorithms.machinery;
 
+import org.neo4j.gds.api.ExportedRelationship;
 import org.neo4j.gds.api.Graph;
 import org.neo4j.gds.api.IdMap;
 import org.neo4j.gds.api.ResultStore;
+import org.neo4j.gds.api.nodeproperties.ValueType;
 import org.neo4j.gds.api.properties.relationships.RelationshipWithPropertyConsumer;
 import org.neo4j.gds.applications.algorithms.metadata.RelationshipsWritten;
 import org.neo4j.gds.core.utils.progress.JobId;
 import org.neo4j.gds.logging.Log;
 
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 public class WriteRelationshipService {
     private final Log log;
@@ -40,22 +44,23 @@ public WriteRelationshipService(Log log, RequestScopedDependencies requestScoped
         this.writeContext = writeContext;
     }
 
-    public RelationshipsWritten write(
+    public RelationshipsWritten writeFromGraph(
         String writeRelationshipType,
         String writeProperty,
-        Graph graph,
+        Graph writeGraph,
         IdMap rootIdMap,
         String taskName,
         Optional<ResultStore> resultStore,
         RelationshipWithPropertyConsumer relationshipWithPropertyConsumer,
         JobId jobId
     ) {
-        return Neo4jDatabaseRelationshipWriter.writeRelationship(
+
+        return Neo4jDatabaseRelationshipWriter.writeRelationshipsFromGraph(
             writeRelationshipType,
             writeProperty,
             requestScopedDependencies.taskRegistryFactory(),
             writeContext.relationshipExporterBuilder(),
-            graph,
+            writeGraph,
             rootIdMap,
             log,
             taskName,
@@ -65,4 +70,33 @@ public RelationshipsWritten write(
             jobId
         );
     }
+
+    public RelationshipsWritten writeFromRelationshipStream(
+        String writeRelationshipType,
+        List<String> properties,
+        List<ValueType> valueTypes,
+        Stream<ExportedRelationship> relationshipStream,
+        IdMap rootIdMap,
+        String taskName,
+        Optional<ResultStore> resultStore,
+        JobId jobId
+    ){
+
+        return Neo4jDatabaseRelationshipWriter.writeRelationshipsFromStream(
+            writeRelationshipType,
+            properties,
+            valueTypes,
+            requestScopedDependencies.taskRegistryFactory(),
+            writeContext.relationshipStreamExporterBuilder(),
+            relationshipStream,
+            rootIdMap,
+            log,
+            taskName,
+            requestScopedDependencies.terminationFlag(),
+            resultStore,
+            jobId
+        );
+    }
+
+
 }