[apache#1464][FOLLOWUP] improvement(spark): print abnormal shuffle se…

…rvers that blocks fail to send
rickyma · Jan 19, 2024 · 6060c21 · 6060c21
1 parent 2bcaf9e
commit 6060c21
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/client-spark/spark2/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java b/client-spark/spark2/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -377,7 +378,9 @@ protected void checkBlockSendResult(Set<Long> blockIds) {
                 + "] failed because "
                 + failedBlockIds.size()
                 + " blocks can't be sent to shuffle server: "
-                + failedBlockIdsWithShuffleServer.values().stream().collect(Collectors.toSet());
+                + failedBlockIdsWithShuffleServer.values().stream()
+                    .flatMap(Collection::stream)
+                    .collect(Collectors.toSet());
         LOG.error(errorMsg);
         throw new RssSendFailedException(errorMsg);
       }

diff --git a/client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java b/client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java
@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -400,7 +401,9 @@ private void checkIfBlocksFailed() {
               + " failed because "
               + failedBlockIds.size()
               + " blocks can't be sent to shuffle server: "
-              + failedBlockIdsWithShuffleServer.values().stream().collect(Collectors.toSet());
+              + failedBlockIdsWithShuffleServer.values().stream()
+                  .flatMap(Collection::stream)
+                  .collect(Collectors.toSet());
       LOG.error(errorMsg);
       throw new RssSendFailedException(errorMsg);
     }