Use node map instead of a throwaway set

Signed-off-by: Daniel Widdis <[email protected]>
opensearch-project · Nov 22, 2023 · d92ee7c · d92ee7c
1 parent e1fee08
commit d92ee7c
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowProcessSorter.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowProcessSorter.java
@@ -29,7 +29,6 @@
 import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -158,20 +157,20 @@ private TimeValue parseTimeout(WorkflowNode node) {
 
     private static List<WorkflowNode> topologicalSort(List<WorkflowNode> workflowNodes, List<WorkflowEdge> workflowEdges) {
         // Basic validation
-        Set<String> nodeIds = new HashSet<>();
+        Map<String, WorkflowNode> nodeMap = new HashMap<>();
         for (WorkflowNode node : workflowNodes) {
-            if (nodeIds.contains(node.id())) {
+            if (nodeMap.containsKey(node.id())) {
                 throw new FlowFrameworkException("Duplicate node id " + node.id() + ".", RestStatus.BAD_REQUEST);
             }
-            nodeIds.add(node.id());
+            nodeMap.put(node.id(), node);
         }
         for (WorkflowEdge edge : workflowEdges) {
             String source = edge.source();
-            if (!nodeIds.contains(source)) {
+            if (!nodeMap.containsKey(source)) {
                 throw new FlowFrameworkException("Edge source " + source + " does not correspond to a node.", RestStatus.BAD_REQUEST);
             }
             String dest = edge.destination();
-            if (!nodeIds.contains(dest)) {
+            if (!nodeMap.containsKey(dest)) {
                 throw new FlowFrameworkException("Edge destination " + dest + " does not correspond to a node.", RestStatus.BAD_REQUEST);
             }
             if (source.equals(dest)) {
@@ -182,7 +181,6 @@ private static List<WorkflowNode> topologicalSort(List<WorkflowNode> workflowNod
         // Build predecessor and successor maps
         Map<WorkflowNode, Set<WorkflowEdge>> predecessorEdges = new HashMap<>();
         Map<WorkflowNode, Set<WorkflowEdge>> successorEdges = new HashMap<>();
-        Map<String, WorkflowNode> nodeMap = workflowNodes.stream().collect(Collectors.toMap(WorkflowNode::id, Function.identity()));
         for (WorkflowEdge edge : workflowEdges) {
             WorkflowNode source = nodeMap.get(edge.source());
             WorkflowNode dest = nodeMap.get(edge.destination());