[SPARK-7076][SPARK-7077][SPARK-7080][SQL] Use managed memory for aggregations

[JoshRosen]

This patch adds managed-memory-based aggregation to Spark SQL / DataFrames. Instead of working with Java objects, this new aggregation path uses sun.misc.Unsafe to manipulate raw memory. This reduces the memory footprint for aggregations, resulting in fewer spills, OutOfMemoryErrors, and garbage collection pauses. As a result, this allows for higher memory utilization. It can also result in better cache locality since objects will be stored closer together in memory.

This feature can be eanbled by setting spark.sql.unsafe.enabled=true. For now, this feature is only supported when codegen is enabled and only supports aggregations for which the grouping columns are primitive numeric types or strings and aggregated values are numeric.

Managing memory with sun.misc.Unsafe

This patch supports both on- and off-heap managed memory.

• In on-heap mode, memory addresses are identified by the combination of a base Object and an offset within that object.
• In off-heap mode, memory is addressed directly with 64-bit long addresses.

To support both modes, functions that manipulate memory accept both baseObject and baseOffset fields. In off-heap mode, we simply pass null as baseObject.

We allocate memory in large chunks, so memory fragmentation and allocation speed are not significant bottlenecks.

By default, we use on-heap mode. To enable off-heap mode, set spark.unsafe.offHeap=true.

To track allocated memory, this patch extends SparkEnv with an ExecutorMemoryManager and supplies each TaskContext with a TaskMemoryManager. These classes work together to track allocations and detect memory leaks.

Compact tuple format

This patch introduces UnsafeRow, a compact row layout. In this format, each tuple has three parts: a null bit set, fixed length values, and variable-length values:

• Rows are always 8-byte word aligned (so their sizes will always be a multiple of 8 bytes)
• The bit set is used for null tracking:
  • Position i is set if and only if field i is null
  • The bit set is aligned to an 8-byte word boundary.
• Every field appears as an 8-byte word in the fixed-length values part:
  • If a field is null, we zero out the values.
  • If a field is variable-length, the word stores a relative offset (w.r.t. the base of the tuple) that points to the beginning of the field's data in the variable-length part.
• Each variable-length data type can have its own encoding:
  • For strings, the first word stores the length of the string and is followed by UTF-8 encoded bytes. If necessary, the end of the string is padded with empty bytes in order to ensure word-alignment.

For example, a tuple that consists 3 fields of type (int, string, string), with value (null, “data”, “bricks”) would look like this:

This format allows us to compare tuples for equality by directly comparing their raw bytes. This also enables fast hashing of tuples.

Hash map for performing aggregations

This patch introduces UnsafeFixedWidthAggregationMap, a hash map for performing aggregations where the aggregation result columns are fixed-with. This map's keys and values are Row objects. UnsafeFixedWidthAggregationMap is implemented on top of BytesToBytesMap, an append-only map which supports byte-array keys and values.

BytesToBytesMap stores pointers to key and value tuples. For each record with a new key, we copy the key and create the aggregation value buffer for that key and put them in a buffer. The hash table then simply stores pointers to the key and value. For each record with an existing key, we simply run the aggregation function to update the values in place.

This map is implemented using open hashing with triangular sequence probing. Each entry stores two words in a long array: the first word stores the address of the key and the second word stores the relative offset from the key tuple to the value tuple, as well as the key's 32-bit hashcode. By storing the full hashcode, we reduce the number of equality checks that need to be performed to handle position collisions ()since the chance of hashcode collision is much lower than position collision).

UnsafeFixedWidthAggregationMap allows regular Spark SQL Row objects to be used when probing the map. Internally, it encodes these rows into UnsafeRow format using UnsafeRowConverter. This conversion has a small overhead that can be eliminated in the future once we use UnsafeRows in other operators.

[SparkQA]

Test build #31245 has finished for PR 5725 at commit eeee512.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds the following public classes (experimental):
  • public final class UnsafeFixedWidthAggregationMap
  • public static class MapEntry
  • public final class UnsafeRow implements MutableRow
  • class UnsafeRowConverter(fieldTypes: Array[DataType])
  • public final class PlatformDependent
  • public class ByteArrayMethods
  • public final class LongArray
  • public final class BitSet
  • public final class BitSetMethods
  • public final class Murmur3_x86_32
  • public final class BytesToBytesMap
  • public final class Location
  • class Doubling implements HashMapGrowthStrategy
  • public class ExecutorMemoryManager
  • public class HeapMemoryAllocator implements MemoryAllocator
  • public class MemoryBlock extends MemoryLocation
  • public class MemoryLocation
  • public final class TaskMemoryManager
  • public class UnsafeMemoryAllocator implements MemoryAllocator
• This patch adds the following new dependencies:
  • spark-unsafe_2.10-1.4.0-SNAPSHOT.jar

[AmplabJenkins]

Merged build finished. Test PASSed.

[AmplabJenkins]

Test PASSed.
Refer to this link for build results (access rights to CI server needed):
https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/31245/
Test PASSed.

[rxin]

Ok merging in master. Let's make sure we address @marmbrus 's comments in followup prs.

[tedyu]

^{sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeFixedWidthAggregationMap.java, line 133 [r14] (raw file):}
Please include writtenLength, unsafeRow.length and javaRow in the assertion so that we have more information when debugging.

---

^{sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/UnsafeRow.java, line 148 [r14] (raw file):}
Condition doesn't match assertion message w.r.t. <=

---

^{unsafe/src/main/java/org/apache/spark/unsafe/array/ByteArrayMethods.java, line 23 [r14] (raw file):}
Declare this class static ?

---

^{unsafe/src/main/java/org/apache/spark/unsafe/array/LongArray.java, line 26 [r14] (raw file):}
nit: in-heap -> on-heap

---

^{unsafe/src/main/java/org/apache/spark/unsafe/bitset/BitSet.java, line 24 [r14] (raw file):}
https://docs.oracle.com/javase/7/docs/api/java/util/BitSet.html doesn't serve your needs ?

---

^{unsafe/src/main/java/org/apache/spark/unsafe/bitset/BitSetMethods.java, line 75 [r14] (raw file):}
Maybe use something similar to http://www.docjar.com/docs/api/sun/misc/Unsafe.html#getInt(Object, long) for speedup ?

---

---

Comments from the review on Reviewable.io

[tgravescs]

I'm not able to build after this went in. Any special requirements to build this?

[info] Compiling 15 Java sources to /home/tgraves/tgravescs_spark/unsafe/target/scala-2.10/classes...
[error] /home/tgraves/tgravescs_spark/unsafe/src/main/java/org/apache/spark/unsafe/PlatformDependent.java:22: warning: sun.misc.Unsafe is Sun proprietary API and may be removed in a future release
[error] import sun.misc.Unsafe;
[error] ^
[error] /home/tgraves/tgravescs_spark/unsafe/src/main/java/org/apache/spark/unsafe/PlatformDependent.java:26: warning: sun.misc.Unsafe is Sun proprietary API and may be removed in a future release
[error] public static final Unsafe UNSAFE;
[error] ^
[error] /home/tgraves/tgravescs_spark/unsafe/src/main/java/org/apache/spark/unsafe/PlatformDependent.java:43: warning: sun.misc.Unsafe is Sun proprietary API and may be removed in a future release
[error] sun.misc.Unsafe unsafe;
[error] ^
[error] /home/tgraves/tgravescs_spark/unsafe/src/main/java/org/apache/spark/unsafe/PlatformDependent.java:45: warning: sun.misc.Unsafe is Sun proprietary API and may be removed in a future release
[error] Field unsafeField = Unsafe.class.getDeclaredField("theUnsafe");
[error] ^
[error] /home/tgraves/tgravescs_spark/unsafe/src/main/java/org/apache/spark/unsafe/PlatformDependent.java:47: warning: sun.misc.Unsafe is Sun proprietary API and may be removed in a future release
[error] unsafe = (sun.misc.Unsafe) unsafeField.get(null);
[error] ^
[error] /home/tgraves/tgravescs_spark/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryBlock.java:53: setMemory(long,long,byte) in sun.misc.Unsafe cannot be applied to (java.lang.Object,long,long,byte)
[error] PlatformDependent.UNSAFE.setMemory(obj, offset, length, (byte) 0);
[error] ^
[error] 1 error
[error] 5 warnings
[error] Compile failed at Apr 30, 2015 2:52:04 PM [0.862s]

[rxin]

Did you have warning turned on into error?

[tgravescs]

not sure what your question is exactly, I think the last one is an error instead of warning. I didn't specifying anything special in build command: mvn -Pyarn -Phadoop-2.4 -DskipTests clean package install

[JoshRosen]

I'm looking into this issue now; these warnings can't be suppressed via the usual annotations mechanism; it requires a special compiler flag. Planning to submit a patch soon; would be great to get your help in testing it.

[tgravescs]

Thanks, Also I should have specified, I'm running on a rhel6 box using jdlk6

[vanzin]

Java 6's output:

$ javap sun.misc.Unsafe | grep setMemory
    public native void setMemory(long, long, byte);

So Spark is using an API that doesn't exist there.

[vanzin]

On a side note, if Spark supports java 6, why are all jenkins builds running java 7?

[JoshRosen]

@tgravescs I've opened #5814 to try to fix the compiler warnings (and error due to Unsafe.setMemory); it would be great if you could test out that commit and see whether it fixes things in your environment.

[JoshRosen]

@vanzin There's an old JIRA for building w/ Java 6 on Jenkins; I've re-opened it: https://issues.apache.org/jira/browse/SPARK-1437

[cloud-fan]

Should this be 8192 * 2^31 * 8 bytes? The length of array is up to Integer.MAX_VALUE.