Optimizations related to zeroing out memory:

- Do not zero out all allocated memory; the zeroing isn't free and in many
  cases it isn't necessary.
- There are some cases where we do want to clear the memory, such as in BitSet.

  It shouldn't be the BitSet object's responsibility to zero out the memory
  block passed to it (since maybe we're passing some memory created by someone
  else and want to interpret it as a bitset). To make the caller's life easier,
  though, I added a MemoryBlock.zero() method for clearing the block.

- In UnsafeGeneratedAggregate, use Arrays.fill to clear the re-used temporary
  row buffer, since this is likely to be much faster than Unsafe.setMemory;
  see http://psy-lob-saw.blogspot.com/2015/04/on-arraysfill-intrinsics-superword-and.html
  for more details.

Author: JoshRosen

sql/core/src/main/scala/org/apache/spark/sql/execution/UnsafeGeneratedAggregate.scala

@@ -17,14 +17,15 @@
 
 package org.apache.spark.sql.execution
 
+import java.util
+
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.trees._
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.plans.physical._
 import org.apache.spark.sql.types._
 import org.apache.spark.unsafe.PlatformDependent
-import org.apache.spark.unsafe.array.ByteArrayMethods
 import org.apache.spark.unsafe.map.BytesToBytesMap
 import org.apache.spark.unsafe.memory.MemoryAllocator
 
@@ -295,8 +296,7 @@ case class UnsafeGeneratedAggregate(
           // Zero out the buffer that's used to hold the current row. This is necessary in order
           // to ensure that rows hash properly, since garbage data from the previous row could
           // otherwise end up as padding in this row.
-          ByteArrayMethods.zeroBytes(
-            unsafeRowBuffer, PlatformDependent.LONG_ARRAY_OFFSET, unsafeRowBuffer.length)
+          util.Arrays.fill(unsafeRowBuffer, 0)
           // Grab the next row from our input iterator and compute its group projection.
           // In the long run, it might be nice to use Unsafe rows for this as well, but for now
           // we'll just rely on the existing code paths to compute the projection.

unsafe/src/main/java/org/apache/spark/unsafe/array/ByteArrayMethods.java

@@ -40,15 +40,6 @@ public static int roundNumberOfBytesToNearestWord(int numBytes) {
     }
   }
 
-  public static void zeroBytes(
-      Object baseObject,
-      long baseOffset,
-      long lengthInBytes) {
-    for (int i = 0; i < lengthInBytes; i++) {
-      PlatformDependent.UNSAFE.putByte(baseObject, baseOffset + i, (byte) 0);
-    }
-  }
-
   /**
    * Optimized  equality check for equal-length byte arrays.
    * @return true if the arrays are equal, false otherwise

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

@@ -22,10 +22,7 @@
 import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.unsafe.bitset.BitSet;
 import org.apache.spark.unsafe.hash.Murmur3_x86_32;
-import org.apache.spark.unsafe.memory.HeapMemoryAllocator;
-import org.apache.spark.unsafe.memory.MemoryAllocator;
-import org.apache.spark.unsafe.memory.MemoryBlock;
-import org.apache.spark.unsafe.memory.MemoryLocation;
+import org.apache.spark.unsafe.memory.*;
 
 import java.lang.IllegalStateException;
 import java.lang.Long;
@@ -452,7 +449,7 @@ public void storeKeyAndValue(
   private void allocate(long capacity) {
     capacity = java.lang.Math.max(nextPowerOf2(capacity), 64);
     longArray = new LongArray(allocator.allocate(capacity * 8 * 2));
-    bitset = new BitSet(allocator.allocate(capacity / 8));
+    bitset = new BitSet(allocator.allocate(capacity / 8).zero());
 
     this.growthThreshold = (long) (capacity * loadFactor);
     this.mask = capacity - 1;
@@ -525,6 +522,9 @@ private void growAndRehash() {
       }
     }
 
+    // TODO: we should probably have a try-finally block here to make sure that we free the allocated
+    // memory even if an error occurs.
+
     // Deallocate the old data structures.
     allocator.free(oldLongArray.memoryBlock());
     allocator.free(oldBitSet.memoryBlock());

unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryAllocator.java

@@ -19,6 +19,10 @@
 
 public interface MemoryAllocator {
 
+  /**
+   * Allocates a contiguous block of memory. Note that the allocated memory is not guaranteed
+   * to be zeroed out (call `zero()` on the result if this is necessary).
+   */
   public MemoryBlock allocate(long size) throws OutOfMemoryError;
 
   public void free(MemoryBlock memory);

unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryBlock.java

@@ -40,6 +40,14 @@ public long size() {
     return length;
   }
 
+  /**
+   * Clear the contents of this memory block.  Returns `this` to facilitate chaining.
+   */
+  public MemoryBlock zero() {
+    PlatformDependent.UNSAFE.setMemory(obj, offset, length, (byte) 0);
+    return this;
+  }
+
   /**
    * Creates a memory block pointing to the memory used by the byte array.
    */

unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java

@@ -27,7 +27,6 @@ public class UnsafeMemoryAllocator implements MemoryAllocator {
   @Override
   public MemoryBlock allocate(long size) throws OutOfMemoryError {
     long address = PlatformDependent.UNSAFE.allocateMemory(size);
-    PlatformDependent.UNSAFE.setMemory(address, size, (byte) 0);
     return new MemoryBlock(null, address, size);
   }
 

unsafe/src/test/java/org/apache/spark/unsafe/bitset/TestBitSet.java

@@ -27,7 +27,7 @@ public class TestBitSet {
 
   private BitSet createBitSet(int capacity) {
     assert capacity % 64 == 0;
-    return new BitSet(MemoryBlock.fromLongArray(new long[capacity / 64]));
+    return new BitSet(MemoryBlock.fromLongArray(new long[capacity / 64]).zero());
   }
 
   @Test