diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/SlotsAllocator.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/SlotsAllocator.java
index a22ef7656..44901668b 100644
--- a/master/src/main/java/org/apache/celeborn/service/deploy/master/SlotsAllocator.java
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/SlotsAllocator.java
@@ -262,12 +262,12 @@ public class SlotsAllocator {
       double fetchTimeWeight) {
     List<List<DiskInfo>> diskGroups = new ArrayList<>();
     usableDisks.sort(
-        (o1, o2) ->
-            Math.toIntExact(
-                (long)
-                    ((o1.avgFlushTime() * flushTimeWeight + o1.avgFetchTime() * fetchTimeWeight)
-                        - (o2.avgFlushTime() * flushTimeWeight
-                            + o2.avgFetchTime() * fetchTimeWeight))));
+        (o1, o2) -> {
+          double delta =
+              (o1.avgFlushTime() * flushTimeWeight + o1.avgFetchTime() * fetchTimeWeight)
+                  - (o2.avgFlushTime() * flushTimeWeight + o2.avgFetchTime() * fetchTimeWeight);
+          return delta < 0 ? -1 : (delta > 0 ? 1 : 0);
+        });
     int diskCount = usableDisks.size();
     int startIndex = 0;
     int groupSizeSize = (int) Math.ceil(usableDisks.size() / (double) diskGroupCount);