[GLUTEN-12464][CORE] Add TaskContext JNI callback for reading Spark task attempt id from native

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})

cpp/core/jni/JniCommon.h

@@ -163,6 +163,10 @@ class JniCommonState {
 
   jmethodID runtimeAwareCtxHandle();
 
+  JavaVM* getJavaVM() const {
+    return vm_;
+  }
+
  private:
   void initialize(JNIEnv* env);
 

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

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

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();
+  }
+}