[CARMEL-640] add max memory limit for result tasks (#45)

Author: wangshisan

core/src/main/scala/org/apache/spark/SparkException.scala

@@ -50,3 +50,10 @@ private[spark] case class ExecutorDeadException(message: String)
 private[spark] class SparkUpgradeException(version: String, message: String, cause: Throwable)
   extends RuntimeException("You may get a different result due to the upgrading of Spark" +
     s" $version: $message", cause)
+
+/**
+ * A poison pill exception, once get a such exception, the scheduler will abort the whole stage
+ * the task belongs to.
+ */
+private[spark] case class PoisonPillException(message: String, cause: Throwable = null)
+  extends SparkException(message, cause)

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -2604,6 +2604,14 @@ object SQLConf {
       .checkValue(_ > 0, "The timeout value must be positive")
       .createWithDefault(10L)
 
+  val TASK_MAX_RETURN_SIZE =
+    buildConf("spark.sql.task.return.size.limit")
+      .doc("This is a limit for result tasks (tasks in the final stage of a query execution)." +
+        " To protect executor from OOM, tasks whose result size over this limit will exist" +
+        " with exception immediately.")
+      .intConf
+      .createWithDefault(1 << 30)
+
   val BROADCAST_HASH_JOIN_OUTPUT_PARTITIONING_EXPAND_LIMIT =
     buildConf("spark.sql.execution.broadcastHashJoin.outputPartitioningExpandLimit")
       .internal()
@@ -3258,6 +3266,8 @@ class SQLConf extends Serializable with Logging {
   def maxNumberForTemporaryTablesPerSession: Long =
     getConf(StaticSQLConf.TEMPORARY_TABLE_MAX_NUM_PER_SESSION)
 
+  def maxTaskResultSize: Int = getConf(TASK_MAX_RETURN_SIZE)
+
   def shuffleAdjustPartitionNumThresholdMin: Double =
     getConf(SHUFFLE_ADJUST_PARTITION_NUM_THRESHOLD_MIN)
 

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

@@ -22,7 +22,7 @@ import java.util.concurrent.atomic.AtomicInteger
 
 import scala.collection.mutable.{ArrayBuffer, ListBuffer}
 
-import org.apache.spark.{broadcast, SparkEnv}
+import org.apache.spark.{broadcast, PoisonPillException, SparkEnv}
 import org.apache.spark.internal.Logging
 import org.apache.spark.io.CompressionCodec
 import org.apache.spark.rdd.{RDD, RDDOperationScope}
@@ -313,6 +313,7 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
    */
   private def getByteArrayRdd(
       n: Int = -1, takeFromEnd: Boolean = false): RDD[(Long, Array[Byte])] = {
+    val sizeLimit = sqlContext.conf.maxTaskResultSize
     execute().mapPartitionsInternal { iter =>
       var count = 0
       val buffer = new Array[Byte](4 << 10)  // 4K
@@ -329,6 +330,11 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
         var i = 0
         count = last.length
         while (i < count) {
+          if (bos.size() > sizeLimit) {
+            throw PoisonPillException(
+              s"Returned too many result, please consider to add a limit clause!" +
+                s" Result size for each task should be under $sizeLimit bytes.")
+          }
           val row = last(i)
           out.writeInt(row.getSizeInBytes)
           row.writeToStream(out, buffer)
@@ -338,6 +344,11 @@ abstract class SparkPlan extends QueryPlan[SparkPlan] with Logging with Serializ
         // `iter.hasNext` may produce one row and buffer it, we should only call it when the
         // limit is not hit.
         while ((n < 0 || count < n) && iter.hasNext) {
+          if (bos.size() > sizeLimit) {
+            throw PoisonPillException(
+              s"Returned too many result, please consider to add a limit clause!" +
+                s" Result size for each task should be under $sizeLimit bytes.")
+          }
           val row = iter.next().asInstanceOf[UnsafeRow]
           out.writeInt(row.getSizeInBytes)
           row.writeToStream(out, buffer)