(incubator-uniffle) branch master updated: [#1807] improvement(server): Optimize judgment of partition whether is huge (#1808)

[zu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zuston pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git


The following commit(s) were added to refs/heads/master by this push:
     new 7aa95ce9d [#1807] improvement(server): Optimize judgment of partition whether is huge (#1808)
7aa95ce9d is described below

commit 7aa95ce9d742b21a02958ca00b58043f6a317a0d
Author: xumanbu <ma...@163.com>
AuthorDate: Tue Jun 18 15:49:14 2024 +0800

    [#1807] improvement(server): Optimize judgment of partition whether is huge (#1808)
    
    ### What changes were proposed in this pull request?
    1、ShuffleTaskInfo add isHugePartition funtion used by hugePartitionTags
    2、ShuffleBuffleManger# isHugePartition utilizes ShuffleTaskInfo#isHugePartition, it may be more efficient and faster.
    
    ### Why are the changes needed?
    
    Fix: #1807
    
    ### Does this PR introduce _any_ user-facing change?
    No.
    
    ### How was this patch tested?
    
    add UT
    
    ---------
    
    Co-authored-by: jam.xu <ja...@vipshop.com>
---
 .../main/java/org/apache/uniffle/server/ShuffleTaskInfo.java |  6 ++++++
 .../apache/uniffle/server/buffer/ShuffleBufferManager.java   |  4 ++--
 .../java/org/apache/uniffle/server/ShuffleTaskInfoTest.java  | 12 ++++++++++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskInfo.java b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskInfo.java
index bbbfded01..e7848963f 100644
--- a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskInfo.java
+++ b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskInfo.java
@@ -200,6 +200,12 @@ public class ShuffleTaskInfo {
     }
   }
 
+  public boolean isHugePartition(int shuffleId, int partitionId) {
+    return existHugePartition.get()
+        && hugePartitionTags.containsKey(shuffleId)
+        && hugePartitionTags.get(shuffleId).contains(partitionId);
+  }
+
   public Set<Integer> getShuffleIds() {
     return partitionDataSizes.keySet();
   }
diff --git a/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java b/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java
index c3506921d..d17716daf 100644
--- a/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java
+++ b/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java
@@ -733,8 +733,8 @@ public class ShuffleBufferManager {
 
   boolean isHugePartition(String appId, int shuffleId, int partitionId) {
     return shuffleTaskManager != null
-        && shuffleTaskManager.getPartitionDataSize(appId, shuffleId, partitionId)
-            > hugePartitionSizeThresholdRef.getSizeAsBytes();
+        && shuffleTaskManager.getShuffleTaskInfo(appId) != null
+        && shuffleTaskManager.getShuffleTaskInfo(appId).isHugePartition(shuffleId, partitionId);
   }
 
   public boolean isHugePartition(long usedPartitionDataSize) {
diff --git a/server/src/test/java/org/apache/uniffle/server/ShuffleTaskInfoTest.java b/server/src/test/java/org/apache/uniffle/server/ShuffleTaskInfoTest.java
index b393943d6..6b5c78d16 100644
--- a/server/src/test/java/org/apache/uniffle/server/ShuffleTaskInfoTest.java
+++ b/server/src/test/java/org/apache/uniffle/server/ShuffleTaskInfoTest.java
@@ -119,6 +119,18 @@ public class ShuffleTaskInfoTest {
     assertEquals(1, ShuffleServerMetrics.counterTotalAppWithHugePartitionNum.get());
   }
 
+  @Test
+  public void isHugePartitionTest() {
+    ShuffleTaskInfo shuffleTaskInfo = new ShuffleTaskInfo("hugePartition_appId");
+    // case1
+    assertFalse(shuffleTaskInfo.hasHugePartition());
+    assertFalse(shuffleTaskInfo.isHugePartition(1, 1));
+    // case2
+    shuffleTaskInfo.markHugePartition(1, 1);
+    assertTrue(shuffleTaskInfo.isHugePartition(1, 1));
+    assertFalse(shuffleTaskInfo.isHugePartition(1, 2));
+  }
+
   @Test
   public void partitionSizeSummaryTest() {
     ShuffleTaskInfo shuffleTaskInfo = new ShuffleTaskInfo("partitionSizeSummaryTest_appId");