This is an automated email from the ASF dual-hosted git repository.

richox pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git


The following commit(s) were added to refs/heads/master by this push:
     new c3144b81 [AURON #1488] Remove isDriverSide method from JniBridge 
(#1489)
c3144b81 is described below

commit c3144b819d8ead78d718417d33b3b47121e8b8ba
Author: zhangmang <[email protected]>
AuthorDate: Wed Oct 22 10:46:20 2025 +0800

    [AURON #1488] Remove isDriverSide method from JniBridge (#1489)
---
 native-engine/auron-jni-bridge/src/jni_bridge.rs                     | 4 ----
 native-engine/datafusion-ext-plans/src/memmgr/spill.rs               | 3 ++-
 .../src/main/java/org/apache/spark/sql/auron/JniBridge.java          | 5 -----
 .../scala/org/apache/spark/sql/auron/memory/OnHeapSpillManager.scala | 4 ++++
 4 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/native-engine/auron-jni-bridge/src/jni_bridge.rs 
b/native-engine/auron-jni-bridge/src/jni_bridge.rs
index 85cb4d0c..05dc8ba3 100644
--- a/native-engine/auron-jni-bridge/src/jni_bridge.rs
+++ b/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",
diff --git a/native-engine/datafusion-ext-plans/src/memmgr/spill.rs 
b/native-engine/datafusion-ext-plans/src/memmgr/spill.rs
index 54f70630..e3b8be07 100644
--- a/native-engine/datafusion-ext-plans/src/memmgr/spill.rs
+++ b/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
diff --git 
a/spark-extension/src/main/java/org/apache/spark/sql/auron/JniBridge.java 
b/spark-extension/src/main/java/org/apache/spark/sql/auron/JniBridge.java
index af035c7a..500f3127 100644
--- a/spark-extension/src/main/java/org/apache/spark/sql/auron/JniBridge.java
+++ b/spark-extension/src/main/java/org/apache/spark/sql/auron/JniBridge.java
@@ -80,11 +80,6 @@ public class JniBridge {
         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
diff --git 
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/OnHeapSpillManager.scala
 
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/OnHeapSpillManager.scala
index 2a898f49..38aa1531 100644
--- 
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/OnHeapSpillManager.scala
+++ 
b/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

Reply via email to