[Improvement] Only report to the shuffle servers that owns the blocks

client/src/main/java/org/apache/uniffle/client/impl/ShuffleWriteClientImpl.java

@@ -520,18 +520,28 @@ public void reportShuffleResult(
     Map<ShuffleServerInfo, List<Integer>> groupedPartitions = Maps.newHashMap();
     Map<Integer, Integer> partitionReportTracker = Maps.newHashMap();
     for (Map.Entry<Integer, List<ShuffleServerInfo>> entry : partitionToServers.entrySet()) {
+      int partitionIdx = entry.getKey();
       for (ShuffleServerInfo ssi : entry.getValue()) {
         if (!groupedPartitions.containsKey(ssi)) {
-          groupedPartitions.putIfAbsent(ssi, Lists.newArrayList());
+          groupedPartitions.put(ssi, Lists.newArrayList());
         }
-        groupedPartitions.get(ssi).add(entry.getKey());
+        groupedPartitions.get(ssi).add(partitionIdx);
+      }
+      if (CollectionUtils.isNotEmpty(partitionToBlockIds.get(partitionIdx))) {
+        partitionReportTracker.putIfAbsent(partitionIdx, 0);
       }
-      partitionReportTracker.putIfAbsent(entry.getKey(), 0);
     }
+
     for (Map.Entry<ShuffleServerInfo, List<Integer>> entry : groupedPartitions.entrySet()) {
       Map<Integer, List<Long>> requestBlockIds = Maps.newHashMap();
       for (Integer partitionId : entry.getValue()) {
-        requestBlockIds.put(partitionId, partitionToBlockIds.get(partitionId));
+        List<Long> blockIds = partitionToBlockIds.get(partitionId);
+        if (CollectionUtils.isNotEmpty(blockIds)) {
+          requestBlockIds.put(partitionId, blockIds);
+        }
+      }
+      if (requestBlockIds.isEmpty()) {
+        continue;
       }
       RssReportShuffleResultRequest request = new RssReportShuffleResultRequest(
           appId, shuffleId, taskAttemptId, requestBlockIds, bitmapNum);
@@ -541,7 +551,7 @@ public void reportShuffleResult(
         if (response.getStatusCode() == ResponseStatusCode.SUCCESS) {
           LOG.info("Report shuffle result to " + ssi + " for appId[" + appId
               + "], shuffleId[" + shuffleId + "] successfully");
-          for (Integer partitionId : entry.getValue()) {
+          for (Integer partitionId : requestBlockIds.keySet()) {
             partitionReportTracker.put(partitionId, partitionReportTracker.get(partitionId) + 1);
           }
         } else {

integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleWithRssClientTest.java

@@ -150,6 +150,57 @@ public void rpcFailTest() throws Exception {
     assertEquals(blockIdBitmap, report);
   }
 
+  @Test
+  public void reportBlocksToShuffleServerIfNecessary() {
+    String testAppId = "reportBlocksToShuffleServerIfNecessary_appId";
+
+    shuffleWriteClientImpl.registerShuffle(
+        shuffleServerInfo1,
+        testAppId,
+        1,
+        Lists.newArrayList(new PartitionRange(1, 1)),
+        new RemoteStorageInfo(""),
+        ShuffleDataDistributionType.NORMAL
+    );
+
+    shuffleWriteClientImpl.registerShuffle(
+        shuffleServerInfo2,
+        testAppId,
+        1,
+        Lists.newArrayList(new PartitionRange(2, 2)),
+        new RemoteStorageInfo(""),
+        ShuffleDataDistributionType.NORMAL
+    );
+
+    Map<Integer, List<ShuffleServerInfo>> partitionToServers = Maps.newHashMap();
+    partitionToServers.put(1, Lists.newArrayList(shuffleServerInfo1));
+    partitionToServers.put(2, Lists.newArrayList(shuffleServerInfo2));
+    Map<Integer, List<Long>> partitionToBlocks = Maps.newHashMap();
+    List<Long> blockIds = Lists.newArrayList();
+
+    int partitionIdx = 1;
+    for (int i = 0; i < 5; i++) {
+      blockIds.add(ClientUtils.getBlockId(partitionIdx, 0, i));
+    }
+    partitionToBlocks.put(partitionIdx, blockIds);
+
+    // case1
+    shuffleWriteClientImpl
+        .reportShuffleResult(partitionToServers, testAppId, 1, 0, partitionToBlocks, 1);
+    Roaring64NavigableMap bitmap = shuffleWriteClientImpl
+        .getShuffleResult("GRPC", Sets.newHashSet(shuffleServerInfo1), testAppId,
+            1, 0);
+    assertTrue(bitmap.isEmpty());
+
+    bitmap = shuffleWriteClientImpl
+        .getShuffleResult("GRPC", Sets.newHashSet(shuffleServerInfo1), testAppId,
+            1, partitionIdx);
+    assertEquals(5, bitmap.getLongCardinality());
+    for (int i = 0; i < 5; i++) {
+      assertTrue(bitmap.contains(partitionToBlocks.get(1).get(i)));
+    }
+  }
+
   @Test
   public void reportMultipleServerTest() throws Exception {
     String testAppId = "reportMultipleServerTest";