[SPARK-59431][CORE] Use ArrayDeque instead of LinkedList for BytesToBytesMap dataPages

core/src/main/java/org/apache/spark/unsafe/map/BytesToBytesMap.java

@@ -20,6 +20,8 @@
 import javax.annotation.Nullable;
 import java.io.File;
 import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Deque;
 import java.util.Iterator;
 import java.util.LinkedList;
 
@@ -103,9 +105,9 @@ public interface KeyOperationsFactory {
   private final TaskMemoryManager taskMemoryManager;
 
   /**
-   * A linked list for tracking all allocated data pages so that we can free all of our memory.
+   * A deque for tracking all allocated data pages so that we can free all of our memory.
    */
-  private final LinkedList<MemoryBlock> dataPages = new LinkedList<>();
+  private final Deque<MemoryBlock> dataPages = new ArrayDeque<>();
 
   /**
    * The data page that will be used to store keys and values for new hashtable entries. When this
@@ -299,6 +301,10 @@ public final class MapIterator implements Iterator<Location> {
     private boolean destructive = false;
     private UnsafeSorterSpillReader reader = null;
 
+    // Used only for non-destructive iteration to walk the data pages in order (a deque has no
+    // random access). The destructive path instead consumes pages from the head of `dataPages`.
+    private Iterator<MemoryBlock> pageIterator = null;
+
     private MapIterator(int numRecords, Location loc, boolean destructive) {
       this.numRecords = numRecords;
       this.loc = loc;
@@ -310,6 +316,8 @@ private MapIterator(int numRecords, Location loc, boolean destructive) {
           freeArray(longArray);
           longArray = null;
         }
+      } else {
+        pageIterator = dataPages.iterator();
       }
     }
 
@@ -323,14 +331,19 @@ private void advanceToNextPage() {
 
       try {
         synchronized (this) {
-          int nextIdx = dataPages.indexOf(currentPage) + 1;
-          if (destructive && currentPage != null) {
-            dataPages.remove(currentPage);
-            pageToFree = currentPage;
-            nextIdx--;
+          final MemoryBlock nextPage;
+          if (destructive) {
+            if (currentPage != null) {
+              assert dataPages.peekFirst() == currentPage;
+              dataPages.removeFirst();
+              pageToFree = currentPage;
+            }
+            nextPage = dataPages.peekFirst();
+          } else {
+            nextPage = pageIterator.hasNext() ? pageIterator.next() : null;
           }
-          if (dataPages.size() > nextIdx) {
-            currentPage = dataPages.get(nextIdx);
+          if (nextPage != null) {
+            currentPage = nextPage;
             pageBaseObject = currentPage.getBaseObject();
             offsetInPage = currentPage.getBaseOffset();
             recordsInPage = UnsafeAlignedOffset.getSize(pageBaseObject, offsetInPage);
@@ -413,7 +426,7 @@ public synchronized long spill(long numBytes) throws IOException {
 
       long released = 0L;
       while (dataPages.size() > 0) {
-        MemoryBlock block = dataPages.getLast();
+        MemoryBlock block = dataPages.peekLast();
         // The currentPage is used, cannot be released
         if (block == currentPage) {
           break;
@@ -435,7 +448,7 @@ public synchronized long spill(long numBytes) throws IOException {
         writer.close();
         spillWriters.add(writer);
 
-        dataPages.removeLast();
+        dataPages.pollLast();
         released += block.size();
         freePage(block);
 
@@ -1030,10 +1043,8 @@ public void free() {
       freeArray(longArray);
       longArray = null;
     }
-    Iterator<MemoryBlock> dataPagesIterator = dataPages.iterator();
-    while (dataPagesIterator.hasNext()) {
-      MemoryBlock dataPage = dataPagesIterator.next();
-      dataPagesIterator.remove();
+    while (!dataPages.isEmpty()) {
+      MemoryBlock dataPage = dataPages.removeFirst();
       freePage(dataPage);
     }
     assert(dataPages.isEmpty());

core/src/test/java/org/apache/spark/unsafe/map/AbstractBytesToBytesMapSuite.java

@@ -350,9 +350,10 @@ public void setAndRetrieveAKey() {
     }
   }
 
-  private void iteratorTestBase(boolean destructive, boolean isWithKeyIndex) throws Exception {
+  private int iteratorTestBase(boolean destructive, boolean isWithKeyIndex, long pageSizeBytes)
+      throws Exception {
     final int size = 4096;
-    BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, size / 2, PAGE_SIZE_BYTES);
+    BytesToBytesMap map = new BytesToBytesMap(taskMemoryManager, size / 2, pageSizeBytes);
     Assertions.assertEquals(size / 2, map.maxNumKeysIndex());
     try {
       for (long i = 0; i < size; i++) {
@@ -390,7 +391,8 @@ private void iteratorTestBase(boolean destructive, boolean isWithKeyIndex) throw
       } else {
         iter = map.iterator();
       }
-      int numPages = map.getNumDataPages();
+      final int totalPages = map.getNumDataPages();
+      int numPages = totalPages;
       int countFreedPages = 0;
       while (iter.hasNext()) {
         final BytesToBytesMap.Location loc = iter.next();
@@ -403,6 +405,8 @@ private void iteratorTestBase(boolean destructive, boolean isWithKeyIndex) throw
           final long key = Platform.getLong(loc.getKeyBase(), loc.getKeyOffset());
           Assertions.assertEquals(value, key);
         }
+        Assertions.assertFalse(valuesSeen.get((int) value),
+          "value " + value + " was seen more than once");
         valuesSeen.set((int) value);
         if (destructive) {
           // The iterator moves onto next page and frees previous page
@@ -419,28 +423,39 @@ private void iteratorTestBase(boolean destructive, boolean isWithKeyIndex) throw
         }
       }
       if (destructive) {
-        // Latest page is not freed by iterator but by map itself
-        Assertions.assertEquals(countFreedPages, numPages - 1);
+        // The iterator frees each page as it advances except the last, which map.free() releases.
+        Assertions.assertEquals(totalPages - 1, countFreedPages);
+        Assertions.assertEquals(1, map.getNumDataPages());
       }
       Assertions.assertEquals(size, valuesSeen.cardinality());
+      return totalPages;
     } finally {
       map.free();
     }
   }
 
   @Test
   public void iteratorTest() throws Exception {
-    iteratorTestBase(false, false);
+    iteratorTestBase(false, false, PAGE_SIZE_BYTES);
   }
 
   @Test
   public void destructiveIteratorTest() throws Exception {
-    iteratorTestBase(true, false);
+    iteratorTestBase(true, false, PAGE_SIZE_BYTES);
+  }
+
+  @Test
+  public void destructiveIteratorManyPagesTest() throws Exception {
+    // A small page size makes records span many data pages, so the destructive iterator
+    // repeatedly removes the head page. The 64 MiB default fits everything in one page.
+    final int totalPages = iteratorTestBase(true, false, 4096);
+    Assertions.assertTrue(totalPages > 1,
+      "expected the records to span multiple data pages, but got " + totalPages);
   }
 
   @Test
   public void iteratorWithKeyIndexTest() throws Exception {
-    iteratorTestBase(false, true);
+    iteratorTestBase(false, true, PAGE_SIZE_BYTES);
   }
 
   @Test