From d82e10d9bf1457bb74a4ac349a117f44fcbeb31a Mon Sep 17 00:00:00 2001
From: Junfan Zhang <zuston@apache.org>
Date: Fri, 17 Mar 2023 15:31:23 +0800
Subject: [PATCH] [#729] improvement: use `foreach` when iterate over
 Roaring64NavigableMap for better performance (#730)

### What changes were proposed in this pull request?

use foreach when iterate over Roaring64NavigableMap for better performance

### Why are the changes needed?

From the doc of https://github.com/RoaringBitmap/RoaringBitmap/issues/44, it will reduce GC footprint and be faster when using the `foreach`

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

No.

### How was this patch tested?

Don't need.
---
 .../common/segment/LocalOrderSegmentSplitter.java        | 6 +-----
 .../org/apache/uniffle/server/ShuffleTaskManager.java    | 9 +++------
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/common/src/main/java/org/apache/uniffle/common/segment/LocalOrderSegmentSplitter.java b/common/src/main/java/org/apache/uniffle/common/segment/LocalOrderSegmentSplitter.java
index d3f912f702..a8bcd0ce72 100644
--- a/common/src/main/java/org/apache/uniffle/common/segment/LocalOrderSegmentSplitter.java
+++ b/common/src/main/java/org/apache/uniffle/common/segment/LocalOrderSegmentSplitter.java
@@ -23,7 +23,6 @@
 import java.util.List;
 
 import com.google.common.collect.Lists;
-import org.roaringbitmap.longlong.LongIterator;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -156,10 +155,7 @@ public List<ShuffleDataSegment> split(ShuffleIndexResult shuffleIndexResult) {
 
   private List<Long> getExpectedTaskIds(Roaring64NavigableMap expectTaskIds) {
     List<Long> taskIds = new ArrayList<>();
-    LongIterator iterator = expectTaskIds.getLongIterator();
-    while (iterator.hasNext()) {
-      taskIds.add(iterator.next());
-    }
+    expectTaskIds.forEach(value -> taskIds.add(value));
     return taskIds;
   }
 }
diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
index e723d22ce4..4095cedf7e 100644
--- a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
+++ b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
@@ -39,7 +39,6 @@
 import com.google.common.collect.Range;
 import com.google.common.collect.Sets;
 import org.apache.commons.collections.CollectionUtils;
-import org.roaringbitmap.longlong.LongIterator;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -441,15 +440,13 @@ public byte[] getFinishedBlockIds(String appId, Integer shuffleId, Set<Integer>
   // filter the specific partition blockId in the bitmap to the resultBitmap
   protected Roaring64NavigableMap getBlockIdsByPartitionId(Set<Integer> requestPartitions,
       Roaring64NavigableMap bitmap, Roaring64NavigableMap resultBitmap) {
-    LongIterator iter = bitmap.getLongIterator();
-    long mask = (1L << Constants.PARTITION_ID_MAX_LENGTH) - 1;
-    while (iter.hasNext()) {
-      long blockId = iter.next();
+    final long mask = (1L << Constants.PARTITION_ID_MAX_LENGTH) - 1;
+    bitmap.forEach(blockId -> {
       int partitionId = Math.toIntExact((blockId >> Constants.TASK_ATTEMPT_ID_MAX_LENGTH) & mask);
       if (requestPartitions.contains(partitionId)) {
         resultBitmap.addLong(blockId);
       }
-    }
+    });
     return resultBitmap;
   }