diff --git a/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java b/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java index 75aefa1247..e8d46a5466 100644 --- a/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java +++ b/merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java @@ -31,12 +31,14 @@ import java.sql.Connection; import java.sql.SQLException; import java.time.Instant; +import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.SortedMap; +import java.util.function.Function; import java.util.stream.Collectors; public final class PostgresResultsCellRepository implements ResultsCellRepository { @@ -424,13 +426,7 @@ private static void postActivities( // Sorts the map by SpanRecord parent ID to ensure foreign key constraints are met. // Entries with null parent IDs are placed first to avoid foreign key violations // for the "span_has_parent_span" constraint. - final var sortedAllActivityRecords = new LinkedHashMap(); - sortedAllActivityRecords.putAll(allActivityRecords.entrySet().stream() - .filter(entry -> entry.getValue().parentId().isEmpty()) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); - sortedAllActivityRecords.putAll(allActivityRecords.entrySet().stream() - .filter(entry -> !entry.getValue().parentId().isEmpty()) - .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))); + final var sortedAllActivityRecords = topoSort(allActivityRecords, $ -> $.parentId().stream().toList()); postActivitiesAction.apply( datasetId, @@ -439,6 +435,29 @@ private static void postActivities( } } + /** + * Take an unsorted map representing a directed acyclic graph and produce a sorted LinkedHashMap where nodes always + * come after their dependencies. The edges must represent an acyclic graph. + * @param nodes a map from keys to values - the keys are used to define dependencies + * @param dependencies - for a given value, what are the keys of its dependencies? + * @return a sorted LinkedHashMap where nodes always come after their dependencies + */ + private static LinkedHashMap topoSort(Map nodes, Function> dependencies) { + final var worklist = new ArrayList<>(nodes.keySet()); + final var output = new LinkedHashMap(); + while (!worklist.isEmpty()) { + for (int i = worklist.size() - 1; i >= 0; i--) { + final var key = worklist.get(i); + final var node = nodes.get(key); + if (dependencies.apply(node).stream().allMatch(output::containsKey)) { + output.put(key, node); + worklist.remove(i); + } + } + } + return output; + } + private static SpanRecord simulatedActivityToRecord(final SimulatedActivity activity) { return new SpanRecord( activity.type(),