fix ut in spark-2.4

apache · Nov 16, 2024 · 063142e · 063142e
1 parent 04bd2a4
commit 063142e
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/.../test/scala/org/apache/celeborn/service/deploy/cluster/PushMergedDataHardSplitSuite.scala b/.../test/scala/org/apache/celeborn/service/deploy/cluster/PushMergedDataHardSplitSuite.scala
@@ -88,11 +88,11 @@ class PushMergedDataHardSplitSuite extends AnyFunSuite
     val partitionLocationMap =
       shuffleClient.getPartitionLocation(SHUFFLE_ID, MAP_NUM, PARTITION_NUM)
     val worker2PartitionIds = mutable.Map.empty[WorkerInfo, ArrayBuffer[Int]]
-    partitionLocationMap.forEach {
-      (partitionId, partitionLocation) =>
-        worker2PartitionIds
-          .getOrElseUpdate(partitionLocation.getWorker, ArrayBuffer.empty)
-          .append(partitionId)
+    for (partitionId <- 0 until PARTITION_NUM) {
+      val partitionLocation = partitionLocationMap.get(partitionId)
+      worker2PartitionIds
+        .getOrElseUpdate(partitionLocation.getWorker, ArrayBuffer.empty)
+        .append(partitionId)
     }
     val partitions = worker2PartitionIds.values.filter(_.size >= 2).head
     assert(partitions.length >= 2)