This is an automated email from the ASF dual-hosted git repository.
taiyang-li pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new cfef50944a [GLUTEN][CORE] Add TaskContext JNI callback for reading
Spark task attempt id from native (#12435)
cfef50944a is described below
commit cfef50944aa7e10db4f6263c5530dff6d32c0dad
Author: 李扬 <[email protected]>
AuthorDate: Tue Jul 7 17:50:15 2026 +0800
[GLUTEN][CORE] Add TaskContext JNI callback for reading Spark task attempt
id from native (#12435)
Instead of extending every JNI entry point (createRuntime / MemoryManager
create/hold/release) to plumb the Spark task attempt id from Java down to
C++ as an extra parameter, expose a small callback surface that native code
uses on demand:
- Java side:
org.apache.gluten.task.TaskContextJniWrapper#currentTaskAttemptId()
reads TaskContext.get().taskAttemptId() on the current thread and returns
-1 when there is no task context.
- C++ side: gluten::getCurrentSparkTaskAttemptId() attaches the current
thread to the JVM as a daemon on demand and calls back into the Java
helper via JNI. The class ref and method id are cached in function-local
statics on first use.
Because Spark's TaskContext is a per-thread ThreadLocal, this returns a
meaningful value whenever the native call is running on an executor task
thread (or any thread inheriting that ThreadLocal), which is exactly when
backends need it.
No signature change to Runtime / MemoryManager / RuntimeJniWrapper /
NativeMemoryManagerJniWrapper. No behavior change for existing backends
(Velox, ClickHouse) that do not query the task attempt id from native.
Change-Id: I3185249796b0c396813dc39f54bd8e8b8589ca2a
Co-authored-by: Aime <[email protected]>
---
cpp/core/CMakeLists.txt | 3 +-
cpp/core/jni/JniCommon.h | 4 +++
cpp/core/jni/TaskContextJniWrapper.cc | 38 ++++++++++++++++++++
cpp/core/jni/TaskContextJniWrapper.h | 33 +++++++++++++++++
.../apache/gluten/task/TaskContextJniWrapper.java | 42 ++++++++++++++++++++++
5 files changed, 119 insertions(+), 1 deletion(-)
diff --git a/cpp/core/CMakeLists.txt b/cpp/core/CMakeLists.txt
index 9530d603f5..f3860166a4 100644
--- a/cpp/core/CMakeLists.txt
+++ b/cpp/core/CMakeLists.txt
@@ -150,7 +150,8 @@ set(SPARK_COLUMNAR_PLUGIN_SRCS
utils/StringUtil.cc
utils/ObjectStore.cc
jni/JniError.cc
- jni/JniCommon.cc)
+ jni/JniCommon.cc
+ jni/TaskContextJniWrapper.cc)
file(MAKE_DIRECTORY ${root_directory}/releases)
add_library(gluten SHARED ${SPARK_COLUMNAR_PLUGIN_SRCS})
diff --git a/cpp/core/jni/JniCommon.h b/cpp/core/jni/JniCommon.h
index 43ddb2226a..bf5a6a746b 100644
--- a/cpp/core/jni/JniCommon.h
+++ b/cpp/core/jni/JniCommon.h
@@ -163,6 +163,10 @@ class JniCommonState {
jmethodID runtimeAwareCtxHandle();
+ JavaVM* getJavaVM() const {
+ return vm_;
+ }
+
private:
void initialize(JNIEnv* env);
diff --git a/cpp/core/jni/TaskContextJniWrapper.cc
b/cpp/core/jni/TaskContextJniWrapper.cc
new file mode 100755
index 0000000000..4349f1cbf1
--- /dev/null
+++ b/cpp/core/jni/TaskContextJniWrapper.cc
@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "jni/TaskContextJniWrapper.h"
+
+#include "jni/JniCommon.h"
+
+namespace gluten {
+
+int64_t getCurrentSparkTaskAttemptId() {
+ JavaVM* vm = getJniCommonState()->getJavaVM();
+ JNIEnv* env = nullptr;
+ attachCurrentThreadAsDaemonOrThrow(vm, &env);
+
+ static jclass taskContextJniWrapperClass =
+ createGlobalClassReferenceOrError(env,
"Lorg/apache/gluten/task/TaskContextJniWrapper;");
+ static jmethodID currentTaskAttemptIdMethod =
+ getStaticMethodIdOrError(env, taskContextJniWrapperClass,
"currentTaskAttemptId", "()J");
+
+ jlong id = env->CallStaticLongMethod(taskContextJniWrapperClass,
currentTaskAttemptIdMethod);
+ checkException(env);
+ return static_cast<int64_t>(id);
+}
+
+} // namespace gluten
diff --git a/cpp/core/jni/TaskContextJniWrapper.h
b/cpp/core/jni/TaskContextJniWrapper.h
new file mode 100755
index 0000000000..87cff54c55
--- /dev/null
+++ b/cpp/core/jni/TaskContextJniWrapper.h
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cstdint>
+
+namespace gluten {
+
+/**
+ * Returns the current thread's Spark task attempt id by calling back into
+ * org.apache.gluten.task.TaskContextJniWrapper#currentTaskAttemptId via JNI.
+ *
+ * Returns -1 when there is no Spark TaskContext on the calling thread (driver,
+ * standalone native worker, etc.). Safe to call from any thread; the helper
+ * attaches the current thread to the JVM as a daemon on demand.
+ */
+int64_t getCurrentSparkTaskAttemptId();
+
+} // namespace gluten
diff --git
a/gluten-arrow/src/main/java/org/apache/gluten/task/TaskContextJniWrapper.java
b/gluten-arrow/src/main/java/org/apache/gluten/task/TaskContextJniWrapper.java
new file mode 100755
index 0000000000..794ee9cb65
--- /dev/null
+++
b/gluten-arrow/src/main/java/org/apache/gluten/task/TaskContextJniWrapper.java
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.task;
+
+import org.apache.spark.TaskContext;
+
+/**
+ * Callback surface that native code uses to obtain the current thread's Spark
task attempt id via
+ * JNI, without having to plumb it explicitly through every JNI entry point.
+ *
+ * <p>Native side reads the value from {@link TaskContext#get()} at the moment
it is needed. Since
+ * {@link TaskContext} is a per-thread ThreadLocal, this only returns a
meaningful value when the
+ * native call is running on the executor task thread (or on a JVM thread
inheriting the same
+ * ThreadLocal). It returns {@code -1L} on the driver or on a native worker
thread that has no
+ * associated Spark task.
+ */
+public final class TaskContextJniWrapper {
+ private TaskContextJniWrapper() {}
+
+ /**
+ * Returns the current thread's Spark task attempt id, or {@code -1L} when
there is no task
+ * context on the calling thread.
+ */
+ public static long currentTaskAttemptId() {
+ TaskContext tc = TaskContext.get();
+ return tc == null ? -1L : tc.taskAttemptId();
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
