Check memory offHeap enabled for CustomResourceProfileExec

wForget · Mar 27, 2024 · 1c7c8cd · 1c7c8cd
1 parent ae9e88f
commit 1c7c8cd
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/...n-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/...n-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
@@ -53,7 +53,12 @@ case class CustomResourceProfileExec(child: SparkPlan) extends UnaryExecNode {
   private val executorMemoryOverhead =
     conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD)
       .getOrElse(sparkContext.getConf.get("spark.executor.memoryOverhead", "1G"))
-  private val executorOffHeapMemory = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY)
+  private val executorOffHeapMemory =
+    if (sparkContext.getConf.getBoolean("spark.memory.offHeap.enabled", false)) {
+      conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY)
+    } else {
+      None
+    }
 
   override lazy val metrics: Map[String, SQLMetric] = {
     val base = Map(

diff --git a/...n-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/...n-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
@@ -53,7 +53,12 @@ case class CustomResourceProfileExec(child: SparkPlan) extends UnaryExecNode {
   private val executorMemoryOverhead =
     conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD)
       .getOrElse(sparkContext.getConf.get("spark.executor.memoryOverhead", "1G"))
-  private val executorOffHeapMemory = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY)
+  private val executorOffHeapMemory =
+    if (sparkContext.getConf.getBoolean("spark.memory.offHeap.enabled", false)) {
+      conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY)
+    } else {
+      None
+    }
 
   override lazy val metrics: Map[String, SQLMetric] = {
     val base = Map(

diff --git a/...n-spark-3-5/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/...n-spark-3-5/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
@@ -53,7 +53,12 @@ case class CustomResourceProfileExec(child: SparkPlan) extends UnaryExecNode {
   private val executorMemoryOverhead =
     conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_MEMORY_OVERHEAD)
       .getOrElse(sparkContext.getConf.get("spark.executor.memoryOverhead", "1G"))
-  private val executorOffHeapMemory = conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY)
+  private val executorOffHeapMemory =
+    if (sparkContext.getConf.getBoolean("spark.memory.offHeap.enabled", false)) {
+      conf.getConf(FINAL_WRITE_STAGE_EXECUTOR_OFF_HEAP_MEMORY)
+    } else {
+      None
+    }
 
   override lazy val metrics: Map[String, SQLMetric] = {
     val base = Map(