[AURON #1488] Remove isDriverSide method from JniBridge

native-engine/auron-jni-bridge/src/jni_bridge.rs

@@ -563,8 +563,6 @@ pub struct JniBridge<'a> {
     pub method_getTaskOnHeapSpillManager_ret: ReturnType,
     pub method_isTaskRunning: JStaticMethodID,
     pub method_isTaskRunning_ret: ReturnType,
-    pub method_isDriverSide: JStaticMethodID,
-    pub method_isDriverSide_ret: ReturnType,
     pub method_openFileAsDataInputWrapper: JStaticMethodID,
     pub method_openFileAsDataInputWrapper_ret: ReturnType,
     pub method_createFileAsDataOutputWrapper: JStaticMethodID,
@@ -626,7 +624,6 @@ impl<'a> JniBridge<'a> {
             method_getTaskOnHeapSpillManager_ret: ReturnType::Object,
             method_isTaskRunning: env.get_static_method_id(class, "isTaskRunning", "()Z")?,
             method_isTaskRunning_ret: ReturnType::Primitive(Primitive::Boolean),
-            method_isDriverSide: env.get_static_method_id(class, "isDriverSide", "()Z")?,
             method_openFileAsDataInputWrapper: env.get_static_method_id(
                 class,
                 "openFileAsDataInputWrapper",
@@ -639,7 +636,6 @@ impl<'a> JniBridge<'a> {
                 "(Lorg/apache/hadoop/fs/FileSystem;Ljava/lang/String;)Lorg/apache/auron/hadoop/fs/FSDataOutputWrapper;",
             )?,
             method_createFileAsDataOutputWrapper_ret: ReturnType::Object,
-            method_isDriverSide_ret: ReturnType::Primitive(Primitive::Boolean),
             method_getDirectMemoryUsed: env.get_static_method_id(
                 class,
                 "getDirectMemoryUsed",

native-engine/datafusion-ext-plans/src/memmgr/spill.rs

@@ -89,7 +89,8 @@ fn spill_compression_codec() -> &'static str {
 }
 
 pub fn try_new_spill(spill_metrics: &SpillMetrics) -> Result<Box<dyn Spill>> {
-    if !is_jni_bridge_inited() || jni_call_static!(JniBridge.isDriverSide() -> bool)? {
+    if !is_jni_bridge_inited() {
+        // is driver
         Ok(Box::new(FileSpill::try_new(spill_metrics)?))
     } else {
         // use on heap spill if on-heap memory is available, otherwise use file spill

spark-extension/src/main/java/org/apache/spark/sql/auron/JniBridge.java

@@ -80,11 +80,6 @@ public static boolean isTaskRunning() {
         return !tc.isCompleted() && !tc.isInterrupted();
     }
 
-    public static boolean isDriverSide() {
-        TaskContext tc = getTaskContext();
-        return tc == null;
-    }
-
     public static FSDataInputWrapper openFileAsDataInputWrapper(FileSystem fs, String path) throws Exception {
         // the path is a URI string, so we need to convert it to a URI object, ref:
         // org.apache.spark.paths.SparkPath.toPath

spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/OnHeapSpillManager.scala

@@ -63,6 +63,10 @@ class OnHeapSpillManager(taskContext: TaskContext)
    */
   @SuppressWarnings(Array("unused"))
   def isOnHeapAvailable: Boolean = {
+    // if driver, tc always null.
+    if (taskContext == null) {
+      return false
+    }
     val memoryPool = OnHeapSpillManagerHelper.getOnHeapExecutionMemoryPool
     val memoryUsed = memoryPool.memoryUsed
     val memoryFree = memoryPool.memoryFree