[hotfix] Refactor OperatorID compartion in StreamGraph JSONGenerator

javabrett · Jan 8, 2020 · efb12fe · efb12fe
1 parent 07ce59e
commit efb12fe
Showing 1 changed file with 8 additions and 17 deletions.
diff --git a/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/JSONGenerator.java b/flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/JSONGenerator.java
@@ -57,23 +57,14 @@ public String getJSON() {
 		ObjectNode json = mapper.createObjectNode();
 		ArrayNode nodes = mapper.createArrayNode();
 		json.put("nodes", nodes);
-		List<Integer> operatorIDs = new ArrayList<Integer>(streamGraph.getVertexIDs());
-		Collections.sort(operatorIDs, new Comparator<Integer>() {
-			@Override
-			public int compare(Integer idOne, Integer idTwo) {
-				boolean isIdOneSinkId = streamGraph.getSinkIDs().contains(idOne);
-				boolean isIdTwoSinkId = streamGraph.getSinkIDs().contains(idTwo);
-				// put sinks at the back
-				if (isIdOneSinkId == isIdTwoSinkId) {
-					return idOne.compareTo(idTwo);
-				} else if (isIdOneSinkId) {
-					return 1;
-				} else {
-					return -1;
-				}
-			}
-		});
-		visit(nodes, operatorIDs, new HashMap<Integer, Integer>());
+
+		List<Integer> operatorIDs = new ArrayList<>(streamGraph.getVertexIDs());
+		Comparator<Integer> operatorIDComparator = Comparator
+			.comparingInt((Integer id) -> streamGraph.getSinkIDs().contains(id) ? 1 : 0)
+			.thenComparingInt(id -> id);
+		operatorIDs.sort(operatorIDComparator);
+
+		visit(nodes, operatorIDs, new HashMap<>());
 
 		return json.toPrettyString();
 	}