diff --git a/algo/src/main/java/org/neo4j/graphalgo/impl/yens/WeightedPathExporter.java b/algo/src/main/java/org/neo4j/graphalgo/impl/yens/WeightedPathExporter.java index 4f6e242a7..81a9cdf5a 100644 --- a/algo/src/main/java/org/neo4j/graphalgo/impl/yens/WeightedPathExporter.java +++ b/algo/src/main/java/org/neo4j/graphalgo/impl/yens/WeightedPathExporter.java @@ -6,6 +6,7 @@ import org.neo4j.graphalgo.core.utils.ParallelUtil; import org.neo4j.graphalgo.core.utils.Pointer; import org.neo4j.graphalgo.core.utils.StatementApi; +import org.neo4j.helpers.collection.Pair; import org.neo4j.internal.kernel.api.exceptions.KernelException; import org.neo4j.kernel.internal.GraphDatabaseAPI; import org.neo4j.values.storable.Values; @@ -13,6 +14,7 @@ import java.util.List; import java.util.concurrent.ExecutorService; import java.util.stream.Collectors; +import java.util.stream.Stream; /** * Specialized exporter for {@link WeightedPath} @@ -96,10 +98,12 @@ private void writeSequential(List paths) { private void writeParallel(List paths) { final Pointer.IntPointer counter = Pointer.wrap(0); - final List tasks = paths.stream() - .sorted(WeightedPath.comparator()) - .map(path -> (Runnable) () -> - export(String.format("%s%d", relPrefix, counter.v++), propertyName, path)) + + Stream> pathsAndRelTypes = paths.stream().sorted(WeightedPath.comparator()) + .map(path -> Pair.of(path, String.format("%s%d", relPrefix, counter.v++))); + + final List tasks = pathsAndRelTypes + .map(pair -> (Runnable) () -> export(pair.other(), propertyName, pair.first())) .collect(Collectors.toList()); ParallelUtil.run(tasks, executorService); }