(auron) branch master updated: [AURON #2107] Fix driver spill NPE error (#2108)

[csy]

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

csy 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 c48e26f6 [AURON #2107] Fix driver spill NPE error (#2108)
c48e26f6 is described below

commit c48e26f6354e22a90e851019262dc0019c02bb48
Author: Yizhong Zhang <zyzzx...@gmail.com>
AuthorDate: Mon Mar 23 10:52:32 2026 +0800

    [AURON #2107] Fix driver spill NPE error (#2108)
    
    # Which issue does this PR close?
    
    Closes #2107
    
    # Rationale for this change
    When `SparkOnHeapSpillManager.current` is called on the driver side,
    `TaskContext.get` returns null, causing a NPE.
    
    # What changes are included in this PR?
    Add a null check when TaskContext is unavailable.
    Update the JNI bridge SIG_TYPE to resolve methods against the
    OnHeapSpillManager.
    
    # Are there any user-facing changes?
    No.
    
    # How was this patch tested?
    UT
    
    Co-authored-by: zhangyizhong <zhangyizhon...@kuaishou.com>
---
 native-engine/auron-jni-bridge/src/jni_bridge.rs                    | 2 +-
 .../org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager.scala | 6 +++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/native-engine/auron-jni-bridge/src/jni_bridge.rs 
b/native-engine/auron-jni-bridge/src/jni_bridge.rs
index a98ae7f5..f9b5d030 100644
--- a/native-engine/auron-jni-bridge/src/jni_bridge.rs
+++ b/native-engine/auron-jni-bridge/src/jni_bridge.rs
@@ -1535,7 +1535,7 @@ pub struct AuronOnHeapSpillManager<'a> {
     pub method_releaseSpill_ret: ReturnType,
 }
 impl<'a> AuronOnHeapSpillManager<'a> {
-    pub const SIG_TYPE: &'static str = 
"org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager";
+    pub const SIG_TYPE: &'static str = 
"org/apache/auron/memory/OnHeapSpillManager";
 
     pub fn new(env: &JNIEnv<'a>) -> JniResult<AuronOnHeapSpillManager<'a>> {
         let class = get_global_jclass(env, Self::SIG_TYPE)?;
diff --git 
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager.scala
 
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager.scala
index b1122502..a69d9ea1 100644
--- 
a/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager.scala
+++ 
b/spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager.scala
@@ -195,6 +195,10 @@ object SparkOnHeapSpillManager extends Logging {
 
   def current: OnHeapSpillManager = {
     val tc = TaskContext.get
-    all.getOrElseUpdate(tc.taskAttemptId(), new SparkOnHeapSpillManager(tc))
+    if (tc != null) {
+      all.getOrElseUpdate(tc.taskAttemptId(), new SparkOnHeapSpillManager(tc))
+    } else {
+      OnHeapSpillManager.getDisabledOnHeapSpillManager
+    }
   }
 }