From 1c7c8cd75d3b39a1ebf536d39db0d1f9bf07c6dc Mon Sep 17 00:00:00 2001
From: wforget <643348094@qq.com>
Date: Wed, 27 Mar 2024 11:52:44 +0800
Subject: [PATCH] Check memory offHeap enabled for CustomResourceProfileExec

---
 .../spark/sql/execution/CustomResourceProfileExec.scala    | 7 ++++++-
 .../spark/sql/execution/CustomResourceProfileExec.scala    | 7 ++++++-
 .../spark/sql/execution/CustomResourceProfileExec.scala    | 7 ++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
index 3698140fbd0..043d6496b22 100644
--- a/extensions/spark/kyuubi-extension-spark-3-3/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
+++ b/extensions/spark/kyuubi-extension-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/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
index 3698140fbd0..043d6496b22 100644
--- a/extensions/spark/kyuubi-extension-spark-3-4/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
+++ b/extensions/spark/kyuubi-extension-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/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala b/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
index 3698140fbd0..043d6496b22 100644
--- a/extensions/spark/kyuubi-extension-spark-3-5/src/main/scala/org/apache/spark/sql/execution/CustomResourceProfileExec.scala
+++ b/extensions/spark/kyuubi-extension-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(