Support vanilla spark 3.1.1

README.md

@@ -99,6 +99,7 @@ You will find jar files under oap-common/target and oap-spark/target.
 
 To enable rdd cache on Intel Optane PMem, you need add the following configurations to `spark-defaults.conf`
 ```
+spark.memory.pmem.extension.enabled true
 spark.memory.pmem.initial.path [Your Optane PMem paths seperate with comma]
 spark.memory.pmem.initial.size [Your Optane PMem size in GB]
 spark.memory.pmem.usable.ratio [from 0 to 1, 0.85 is recommended]

pom.xml

@@ -15,7 +15,7 @@
         <junit.version>4.12</junit.version>
         <scala.version>2.12.10</scala.version>
         <scala.binary.version>2.12</scala.binary.version>
-        <spark.internal.version>3.0.0</spark.internal.version>
+        <spark.internal.version>3.1.1</spark.internal.version>
         <maven-patch-plugin.version>1.2</maven-patch-plugin.version>
         <maven-resource-plugin.version>2.6</maven-resource-plugin.version>
         <jetty.version>9.4.39.v20210325</jetty.version>

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

@@ -393,10 +393,12 @@ public void remove() {
     }
 
     private void handleFailedDelete() {
-      // remove the spill file from disk
-      File file = spillWriters.removeFirst().getFile();
-      if (file != null && file.exists() && !file.delete()) {
-        logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+      if (spillWriters.size() > 0) {
+        // remove the spill file from disk
+        File file = spillWriters.removeFirst().getFile();
+        if (file != null && file.exists() && !file.delete()) {
+          logger.error("Was unable to delete spill file {}", file.getAbsolutePath());
+        }
       }
     }
   }
@@ -406,17 +408,10 @@ private void handleFailedDelete() {
    *
    * For efficiency, all calls to `next()` will return the same {@link Location} object.
    *
-   * If any other lookups or operations are performed on this map while iterating over it, including
-   * `lookup()`, the behavior of the returned iterator is undefined.
+   * The returned iterator is thread-safe. However if the map is modified while iterating over it,
+   * the behavior of the returned iterator is undefined.
    */
   public MapIterator iterator() {
-    return new MapIterator(numValues, loc, false);
-  }
-
-  /**
-   * Returns a thread safe iterator that iterates of the entries of this map.
-   */
-  public MapIterator safeIterator() {
     return new MapIterator(numValues, new Location(), false);
   }
 
@@ -427,19 +422,82 @@ public MapIterator safeIterator() {
    *
    * For efficiency, all calls to `next()` will return the same {@link Location} object.
    *
-   * If any other lookups or operations are performed on this map while iterating over it, including
-   * `lookup()`, the behavior of the returned iterator is undefined.
+   * The returned iterator is thread-safe. However if the map is modified while iterating over it,
+   * the behavior of the returned iterator is undefined.
    */
   public MapIterator destructiveIterator() {
     updatePeakMemoryUsed();
-    return new MapIterator(numValues, loc, true);
+    return new MapIterator(numValues, new Location(), true);
+  }
+
+  /**
+   * Iterator for the entries of this map. This is to first iterate over key indices in
+   * `longArray` then accessing values in `dataPages`. NOTE: this is different from `MapIterator`
+   * in the sense that key index is preserved here
+   * (See `UnsafeHashedRelation` for example of usage).
+   */
+  public final class MapIteratorWithKeyIndex implements Iterator<Location> {
+
+    /**
+     * The index in `longArray` where the key is stored.
+     */
+    private int keyIndex = 0;
+
+    private int numRecords;
+    private final Location loc;
+
+    private MapIteratorWithKeyIndex() {
+      this.numRecords = numValues;
+      this.loc = new Location();
+    }
+
+    @Override
+    public boolean hasNext() {
+      return numRecords > 0;
+    }
+
+    @Override
+    public Location next() {
+      if (!loc.isDefined() || !loc.nextValue()) {
+        while (longArray.get(keyIndex * 2) == 0) {
+          keyIndex++;
+        }
+        loc.with(keyIndex, 0, true);
+        keyIndex++;
+      }
+      numRecords--;
+      return loc;
+    }
+  }
+
+  /**
+   * Returns an iterator for iterating over the entries of this map,
+   * by first iterating over the key index inside hash map's `longArray`.
+   *
+   * For efficiency, all calls to `next()` will return the same {@link Location} object.
+   *
+   * The returned iterator is NOT thread-safe. If the map is modified while iterating over it,
+   * the behavior of the returned iterator is undefined.
+   */
+  public MapIteratorWithKeyIndex iteratorWithKeyIndex() {
+    return new MapIteratorWithKeyIndex();
+  }
+
+  /**
+   * The maximum number of allowed keys index.
+   *
+   * The value of allowed keys index is in the range of [0, maxNumKeysIndex - 1].
+   */
+  public int maxNumKeysIndex() {
+    return (int) (longArray.size() / 2);
   }
 
   /**
    * Looks up a key, and return a {@link Location} handle that can be used to test existence
    * and read/write values.
    *
-   * This function always return the same {@link Location} instance to avoid object allocation.
+   * This function always returns the same {@link Location} instance to avoid object allocation.
+   * This function is not thread-safe.
    */
   public Location lookup(Object keyBase, long keyOffset, int keyLength) {
     safeLookup(keyBase, keyOffset, keyLength, loc,
@@ -451,7 +509,8 @@ public Location lookup(Object keyBase, long keyOffset, int keyLength) {
    * Looks up a key, and return a {@link Location} handle that can be used to test existence
    * and read/write values.
    *
-   * This function always return the same {@link Location} instance to avoid object allocation.
+   * This function always returns the same {@link Location} instance to avoid object allocation.
+   * This function is not thread-safe.
    */
   public Location lookup(Object keyBase, long keyOffset, int keyLength, int hash) {
     safeLookup(keyBase, keyOffset, keyLength, loc, hash);
@@ -606,6 +665,14 @@ public boolean isDefined() {
       return isDefined;
     }
 
+    /**
+     * Returns index for key.
+     */
+    public int getKeyIndex() {
+      assert (isDefined);
+      return pos;
+    }
+
     /**
      * Returns the base object for key.
      */
@@ -743,14 +810,23 @@ public boolean append(Object kbase, long koff, int klen, Object vbase, long voff
         longArray.set(pos * 2 + 1, keyHashcode);
         isDefined = true;
 
+        // If the map has reached its growth threshold, try to grow it.
+        if (numKeys >= growthThreshold) {
         // We use two array entries per key, so the array size is twice the capacity.
         // We should compare the current capacity of the array, instead of its size.
-        if (numKeys >= growthThreshold && longArray.size() / 2 < MAX_CAPACITY) {
+        if (longArray.size() / 2 < MAX_CAPACITY) {
           try {
             growAndRehash();
           } catch (SparkOutOfMemoryError oom) {
             canGrowArray = false;
           }
+          } else {
+            // The map is already at MAX_CAPACITY and cannot grow. Instead, we prevent it from
+            // accepting any more new elements to make sure we don't exceed the load factor. If we
+            // need to spill later, this allows UnsafeKVExternalSorter to reuse the array for
+            // sorting.
+            canGrowArray = false;
+          }
         }
       }
       return true;

src/main/java/org/apache/spark/util/collection/unsafe/sort/PMemReader.java

@@ -79,6 +79,11 @@ public int getRecordLength() {
     public long getKeyPrefix() {
         return keyPrefix;
     }
+
+   @Override
+    public long getCurrentPageNumber() {
+        throw new UnsupportedOperationException();
+   } 
 
     @Override
     public void close() {

src/main/java/org/apache/spark/util/collection/unsafe/sort/PMemReaderForUnsafeExternalSorter.java

@@ -64,6 +64,11 @@ public void loadNext() {
     public int getNumRecords() {
         return numRecords;
     }
+
+    @Override
+    public long getCurrentPageNumber() {
+        throw new UnsupportedOperationException();
+    } 
 
     /**
      * load more PMem records in the buffer

src/main/java/org/apache/spark/util/collection/unsafe/sort/SortedPMemPageSpillWriter.java

@@ -212,6 +212,12 @@ public SortedPMemPageSpillReader() throws IOException{
                 numRecordsOnDisk = diskSpillWriter.recordsSpilled();
             }
         }
+
+        @Override
+        public long getCurrentPageNumber() {
+           throw new UnsupportedOperationException();
+        }
+
         @Override
         public boolean hasNext() {
             return curNumOfRec < numRecordsOnPMem + numRecordsOnDisk;