[AURON #2107] Fix driver spill NPE error#2108
Conversation
There was a problem hiding this comment.
Pull request overview
This PR fixes a driver-side null dereference when resolving the current on-heap spill manager in Spark, and aligns the JNI bridge to resolve spill methods against the shared OnHeapSpillManager interface so calls work for all implementations (including the disabled manager).
Changes:
- Add a null check in
SparkOnHeapSpillManager.currentto avoid NPE whenTaskContext.getis unavailable (driver-side). - Update the JNI bridge class signature for spill manager method resolution to
org/apache/auron/memory/OnHeapSpillManager.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| spark-extension/src/main/scala/org/apache/spark/sql/auron/memory/SparkOnHeapSpillManager.scala | Prevents driver-side NPE by returning a disabled spill manager when TaskContext is null. |
| native-engine/auron-jni-bridge/src/jni_bridge.rs | Resolves JNI method IDs against the OnHeapSpillManager interface for compatibility across implementations. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if (tc != null) { | ||
| all.getOrElseUpdate(tc.taskAttemptId(), new SparkOnHeapSpillManager(tc)) | ||
| } else { | ||
| OnHeapSpillManager.getDisabledOnHeapSpillManager | ||
| } |
There was a problem hiding this comment.
OnHeapSpillManager.getDisabledOnHeapSpillManager() creates a new anonymous instance on each call (see auron-core/.../OnHeapSpillManager.java). Since current may be called repeatedly on the driver (where TaskContext.get == null), consider caching a single disabled manager instance in SparkOnHeapSpillManager (e.g., a private val) and returning that to avoid unnecessary allocations/GC churn.
Which issue does this PR close?
Closes #2107
Rationale for this change
When
SparkOnHeapSpillManager.currentis called on the driver side,TaskContext.getreturns 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