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