Add configuration options for heap vs. offheap

Author: JoshRosen

sql/core/src/main/scala/org/apache/spark/sql/SQLConf.scala

@@ -30,7 +30,8 @@ private[spark] object SQLConf {
   val DEFAULT_SIZE_IN_BYTES = "spark.sql.defaultSizeInBytes"
   val SHUFFLE_PARTITIONS = "spark.sql.shuffle.partitions"
   val CODEGEN_ENABLED = "spark.sql.codegen"
-  val UNSAFE_ENABLED = "spark.sql.unsafe"
+  val UNSAFE_ENABLED = "spark.sql.unsafe.enabled"
+  val UNSAFE_USE_OFF_HEAP = "spark.sql.unsafe.offHeap"
   val DIALECT = "spark.sql.dialect"
 
   val PARQUET_BINARY_AS_STRING = "spark.sql.parquet.binaryAsString"
@@ -150,8 +151,21 @@ private[sql] class SQLConf extends Serializable {
    */
   private[spark] def codegenEnabled: Boolean = getConf(CODEGEN_ENABLED, "false").toBoolean
 
+  /**
+   * When set to true, Spark SQL will use managed memory for certain operations.  This option only
+   * takes effect if codegen is enabled.
+   *
+   * Defaults to false as this feature is currently experimental.
+   */
   private[spark] def unsafeEnabled: Boolean = getConf(UNSAFE_ENABLED, "false").toBoolean
 
+  /**
+   * When set to true, Spark SQL will use off-heap memory allocation for managed memory operations.
+   *
+   * Defaults to false.
+   */
+  private[spark] def unsafeUseOffHeap: Boolean = getConf(UNSAFE_USE_OFF_HEAP, "false").toBoolean
+
   private[spark] def useSqlSerializer2: Boolean = getConf(USE_SQL_SERIALIZER2, "true").toBoolean
 
   /**

sql/core/src/main/scala/org/apache/spark/sql/SQLContext.scala

@@ -1013,6 +1013,8 @@ class SQLContext(@transient val sparkContext: SparkContext)
 
     def unsafeEnabled: Boolean = self.conf.unsafeEnabled
 
+    def unsafeUseOffHeap: Boolean = self.conf.unsafeUseOffHeap
+
     def numPartitions: Int = self.conf.numShufflePartitions
 
     def strategies: Seq[Strategy] =

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

@@ -43,14 +43,16 @@ case class AggregateEvaluation(
  * @param aggregateExpressions expressions that are computed for each group.
  * @param child the input data source.
  * @param unsafeEnabled whether to allow Unsafe-based aggregation buffers to be used.
+ * @param useOffHeap whether to use off-heap allocation (only takes effect if unsafeEnabled=true)
  */
 @DeveloperApi
 case class GeneratedAggregate(
     partial: Boolean,
     groupingExpressions: Seq[Expression],
     aggregateExpressions: Seq[NamedExpression],
     child: SparkPlan,
-    unsafeEnabled: Boolean)
+    unsafeEnabled: Boolean,
+    useOffHeap: Boolean)
   extends UnaryNode {
 
   override def requiredChildDistribution: Seq[Distribution] =
@@ -289,7 +291,7 @@ case class GeneratedAggregate(
           newAggregationBuffer(EmptyRow),
           aggregationBufferSchema,
           groupKeySchema,
-          MemoryAllocator.UNSAFE,
+          if (useOffHeap) MemoryAllocator.UNSAFE else MemoryAllocator.HEAP,
           1024 * 16,
           false
         )

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

@@ -141,8 +141,10 @@ private[sql] abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
               groupingExpressions,
               partialComputation,
               planLater(child),
-              unsafeEnabled),
-          unsafeEnabled) :: Nil
+              unsafeEnabled,
+              unsafeUseOffHeap),
+          unsafeEnabled,
+          unsafeUseOffHeap) :: Nil
 
       // Cases where some aggregate can not be codegened
       case PartialAggregation(

sql/core/src/test/scala/org/apache/spark/sql/SQLQuerySuite.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql
 
 import org.scalatest.BeforeAndAfterAll
 
-import org.apache.spark.sql.execution.{GeneratedAggregate}
+import org.apache.spark.sql.execution.GeneratedAggregate
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.TestData._
 import org.apache.spark.sql.test.TestSQLContext