[Improvement] Refactor getPartitionRange to calculate range directly (#…

…447)

### What changes were proposed in this pull request?
 1. Refactor ShuffleStorageUtils.getPartitionRange to calculate range directly.
 2. The runtime of the function will be O(1) as opposed to O(n) in the original implementation.

### Why are the changes needed?
fix #443
The current implementation of ShuffleStorageUtils.getPartitionRange is O(n) which is not performant and can be done in O(1) by directly computing it using the partitionId.

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
existing UTs + new UTs

storage/src/main/java/org/apache/uniffle/storage/util/ShuffleStorageUtils.java

@@ -167,15 +167,12 @@ public static String getShuffleDataPathWithRange(
 
   public static int[] getPartitionRange(int partitionId, int partitionNumPerRange, int partitionNum) {
     int[] range = null;
-    int prNum = partitionNum % partitionNumPerRange == 0
-        ? partitionNum / partitionNumPerRange : partitionNum / partitionNumPerRange + 1;
-    for (int i = 0; i < prNum; i++) {
-      int start = i * partitionNumPerRange;
-      int end = (i + 1) * partitionNumPerRange - 1;
-      if (partitionId >= start && partitionId <= end) {
-        range = new int[]{start, end};
-        break;
-      }
+    int prNum = partitionId / partitionNumPerRange;
+    if (partitionId < 0 || partitionId >= partitionNum) {
+      LOG.warn("Invalid partitionId. partitionId:{} ,partitionNumPerRange: {}, partitionNum: {}",
+              partitionId, partitionNumPerRange, partitionNum);
+    } else {
+      range = new int[]{partitionNumPerRange * prNum, partitionNumPerRange * (prNum + 1) - 1};
     }
     return range;
   }

storage/src/test/java/org/apache/uniffle/storage/util/ShuffleStorageUtilsTest.java

@@ -251,6 +251,8 @@ public void getPartitionRangeTest() {
     int[] range = ShuffleStorageUtils.getPartitionRange(0, 1, 5);
     assertEquals(0, range[0]);
     assertEquals(0, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(5, 1, 5);
+    assertEquals(null, range);
     range = ShuffleStorageUtils.getPartitionRange(0, 2, 5);
     assertEquals(0, range[0]);
     assertEquals(1, range[1]);
@@ -266,5 +268,50 @@ public void getPartitionRangeTest() {
     range = ShuffleStorageUtils.getPartitionRange(2, 3, 5);
     assertEquals(0, range[0]);
     assertEquals(2, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(-1, 3, 5);
+    assertEquals(null, range);
+    range = ShuffleStorageUtils.getPartitionRange(1, 3, 0);
+    assertEquals(null, range);
+    range = ShuffleStorageUtils.getPartitionRange(4, 2, 3);
+    assertEquals(null, range);
+    range = ShuffleStorageUtils.getPartitionRange(4, 2, 5);
+    assertEquals(4, range[0]);
+    assertEquals(5, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(3, 2, 2);
+    assertEquals(null, range);
+    range = ShuffleStorageUtils.getPartitionRange(0, 2, 2);
+    assertEquals(0, range[0]);
+    assertEquals(1, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(3, 2, 4);
+    assertEquals(2, range[0]);
+    assertEquals(3, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(4, 3, 5);
+    assertEquals(3, range[0]);
+    assertEquals(5, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(5, 3, 5);
+    assertEquals(null, range);
+    range = ShuffleStorageUtils.getPartitionRange(0, 3, 5);
+    assertEquals(0, range[0]);
+    assertEquals(2, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(9, 3, 11);
+    assertEquals(9, range[0]);
+    assertEquals(11, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(9, 3, 12);
+    assertEquals(9, range[0]);
+    assertEquals(11, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(0, 7, 12);
+    assertEquals(0, range[0]);
+    assertEquals(6, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(3, 7, 12);
+    assertEquals(0, range[0]);
+    assertEquals(6, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(7, 7, 12);
+    assertEquals(7, range[0]);
+    assertEquals(13, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(8, 7, 12);
+    assertEquals(7, range[0]);
+    assertEquals(13, range[1]);
+    range = ShuffleStorageUtils.getPartitionRange(12, 7, 12);
+    assertEquals(null, range);
   }
 }