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 26f04d815c [GLUTEN-12535][CORE] Support backend-specific JNI input
adapters (#12535)
26f04d815c is described below
commit 26f04d815c28c4253fe693a02c3b69f4b385ba58
Author: wangxinshuo-bolt <[email protected]>
AuthorDate: Fri Jul 24 14:10:04 2026 +0800
[GLUTEN-12535][CORE] Support backend-specific JNI input adapters (#12535)
* [CORE] Add reusable JNI input adapters for backend runtimes
* Remove useless test
fix
* fix
* [CORE] Decouple JNI input adapters from Runtime
* [CORE] Move shuffle stream JNI bridge to JniCommon
---
cpp/core/jni/JniCommon.cc | 180 +++++++++++++++++++++++++++++++++++++++++++++
cpp/core/jni/JniCommon.h | 39 ++++++++++
cpp/core/jni/JniWrapper.cc | 131 +--------------------------------
3 files changed, 220 insertions(+), 130 deletions(-)
diff --git a/cpp/core/jni/JniCommon.cc b/cpp/core/jni/JniCommon.cc
index 76a07aa2d8..1f57e25a17 100644
--- a/cpp/core/jni/JniCommon.cc
+++ b/cpp/core/jni/JniCommon.cc
@@ -18,6 +18,99 @@
#include "JniCommon.h"
#include <folly/system/ThreadName.h>
+#include "utils/ArrowStatus.h"
+
+namespace {
+
+std::unordered_map<std::string, gluten::JniInputIteratorFactory>&
jniInputIteratorFactories() {
+ static std::unordered_map<std::string, gluten::JniInputIteratorFactory>
factories;
+ return factories;
+}
+
+std::mutex& jniInputIteratorFactoriesMutex() {
+ static std::mutex mutex;
+ return mutex;
+}
+
+class JavaInputStreamAdaptor final : public arrow::io::InputStream {
+ public:
+ JavaInputStreamAdaptor(JNIEnv* env, arrow::MemoryPool* pool, jobject jniIn)
: pool_(pool) {
+ // IMPORTANT: DO NOT USE LOCAL REF IN DIFFERENT THREAD
+ if (env->GetJavaVM(&vm_) != JNI_OK) {
+ std::string errorMessage = "Unable to get JavaVM instance";
+ throw gluten::GlutenException(errorMessage);
+ }
+ jniIn_ = env->NewGlobalRef(jniIn);
+ }
+
+ ~JavaInputStreamAdaptor() override {
+ try {
+ auto status = JavaInputStreamAdaptor::Close();
+ if (!status.ok()) {
+ LOG(WARNING) << __func__ << " call JavaInputStreamAdaptor::Close()
failed, status:" << status.ToString();
+ }
+ } catch (std::exception& e) {
+ LOG(WARNING) << __func__ << " call JavaInputStreamAdaptor::Close() got
exception:" << e.what();
+ }
+ }
+
+ // not thread safe
+ arrow::Status Close() override {
+ if (closed_) {
+ return arrow::Status::OK();
+ }
+ JNIEnv* env;
+ attachCurrentThreadAsDaemonOrThrow(vm_, &env);
+ env->CallVoidMethod(jniIn_,
gluten::getJniCommonState()->jniByteInputStreamClose());
+ checkException(env);
+ env->DeleteGlobalRef(jniIn_);
+ // Do NOT call DetachCurrentThread() here.
+ // libhdfs.so caches JNIEnv* in thread-local storage after
AttachCurrentThread.
+ // If we detach, libhdfs's TLS cache becomes stale — the next HDFS call via
+ // libhdfs returns the stale env, causing SIGSEGV in jni_NewStringUTF.
+ // Daemon-attached threads are safe to leave attached; they won't block
JVM shutdown.
+ closed_ = true;
+ return arrow::Status::OK();
+ }
+
+ arrow::Result<int64_t> Tell() const override {
+ JNIEnv* env;
+ attachCurrentThreadAsDaemonOrThrow(vm_, &env);
+ jlong told = env->CallLongMethod(jniIn_,
gluten::getJniCommonState()->jniByteInputStreamTell());
+ checkException(env);
+ return told;
+ }
+
+ bool closed() const override {
+ return closed_;
+ }
+
+ arrow::Result<int64_t> Read(int64_t nbytes, void* out) override {
+ JNIEnv* env;
+ attachCurrentThreadAsDaemonOrThrow(vm_, &env);
+ jlong read = env->CallLongMethod(
+ jniIn_, gluten::getJniCommonState()->jniByteInputStreamRead(),
reinterpret_cast<jlong>(out), nbytes);
+ checkException(env);
+ return read;
+ }
+
+ arrow::Result<std::shared_ptr<arrow::Buffer>> Read(int64_t nbytes) override {
+ GLUTEN_ASSIGN_OR_THROW(auto buffer, arrow::AllocateResizableBuffer(nbytes,
pool_))
+ GLUTEN_ASSIGN_OR_THROW(int64_t bytes_read, Read(nbytes,
buffer->mutable_data()))
+ GLUTEN_THROW_NOT_OK(buffer->Resize(bytes_read, false));
+ buffer->ZeroPadding();
+ return std::move(buffer);
+ }
+
+ private:
+ arrow::MemoryPool* pool_;
+ JavaVM* vm_;
+ jobject jniIn_;
+ bool closed_ = false;
+};
+
+} // namespace
+
void gluten::JniCommonState::ensureInitialized(JNIEnv* env) {
std::lock_guard<std::mutex> lockGuard(mtx_);
if (initialized_) {
@@ -38,9 +131,41 @@ jmethodID gluten::JniCommonState::runtimeAwareCtxHandle() {
return runtimeAwareCtxHandle_;
}
+jmethodID gluten::JniCommonState::jniByteInputStreamRead() {
+ assertInitialized();
+ return jniByteInputStreamRead_;
+}
+
+jmethodID gluten::JniCommonState::jniByteInputStreamTell() {
+ assertInitialized();
+ return jniByteInputStreamTell_;
+}
+
+jmethodID gluten::JniCommonState::jniByteInputStreamClose() {
+ assertInitialized();
+ return jniByteInputStreamClose_;
+}
+
+jmethodID gluten::JniCommonState::shuffleStreamReaderNextStream() {
+ assertInitialized();
+ return shuffleStreamReaderNextStream_;
+}
+
void gluten::JniCommonState::initialize(JNIEnv* env) {
runtimeAwareClass_ = createGlobalClassReference(env,
"Lorg/apache/gluten/runtime/RuntimeAware;");
runtimeAwareCtxHandle_ = getMethodIdOrError(env, runtimeAwareClass_,
"rtHandle", "()J");
+
+ jniByteInputStreamClass_ =
+ createGlobalClassReferenceOrError(env,
"Lorg/apache/gluten/vectorized/JniByteInputStream;");
+ jniByteInputStreamRead_ = getMethodIdOrError(env, jniByteInputStreamClass_,
"read", "(JJ)J");
+ jniByteInputStreamTell_ = getMethodIdOrError(env, jniByteInputStreamClass_,
"tell", "()J");
+ jniByteInputStreamClose_ = getMethodIdOrError(env, jniByteInputStreamClass_,
"close", "()V");
+
+ shuffleStreamReaderClass_ =
+ createGlobalClassReferenceOrError(env,
"Lorg/apache/gluten/vectorized/ShuffleStreamReader;");
+ shuffleStreamReaderNextStream_ = getMethodIdOrError(
+ env, shuffleStreamReaderClass_, "nextStream",
"()Lorg/apache/gluten/vectorized/JniByteInputStream;");
+
JavaVM* vm;
if (env->GetJavaVM(&vm) != JNI_OK) {
throw gluten::GlutenException("Unable to get JavaVM instance");
@@ -56,6 +181,8 @@ void gluten::JniCommonState::close() {
JNIEnv* env = nullptr;
attachCurrentThreadAsDaemonOrThrow(vm_, &env);
env->DeleteGlobalRef(runtimeAwareClass_);
+ env->DeleteGlobalRef(jniByteInputStreamClass_);
+ env->DeleteGlobalRef(shuffleStreamReaderClass_);
closed_ = true;
}
@@ -67,6 +194,59 @@ gluten::Runtime* gluten::getRuntime(JNIEnv* env, jobject
runtimeAware) {
return ctx;
}
+void gluten::registerJniInputIteratorFactory(const std::string& runtimeKind,
JniInputIteratorFactory factory) {
+ GLUTEN_CHECK(!runtimeKind.empty(), "JNI input iterator factory runtime kind
must not be empty");
+ GLUTEN_CHECK(static_cast<bool>(factory), "JNI input iterator factory must
not be empty");
+
+ std::lock_guard<std::mutex> lock(jniInputIteratorFactoriesMutex());
+ const bool inserted = jniInputIteratorFactories().emplace(runtimeKind,
std::move(factory)).second;
+ GLUTEN_CHECK(inserted, "JNI input iterator factory already registered for "
+ runtimeKind);
+}
+
+std::unique_ptr<gluten::ColumnarBatchIterator>
+gluten::createJniInputIterator(JNIEnv* env, jobject iterator, Runtime*
runtime, int32_t iteratorIndex) {
+ GLUTEN_CHECK(runtime != nullptr, "Runtime must not be null");
+
+ JniInputIteratorFactory factory;
+ {
+ std::lock_guard<std::mutex> lock(jniInputIteratorFactoriesMutex());
+ const auto it = jniInputIteratorFactories().find(runtime->kind());
+ if (it != jniInputIteratorFactories().end()) {
+ factory = it->second;
+ }
+ }
+
+ if (factory) {
+ return factory(env, iterator, runtime, iteratorIndex);
+ }
+ return std::make_unique<JniColumnarBatchIterator>(env, iterator, runtime,
iteratorIndex);
+}
+
+gluten::ShuffleStreamReader::ShuffleStreamReader(JNIEnv* env, jobject reader) {
+ if (env->GetJavaVM(&vm_) != JNI_OK) {
+ throw GlutenException("Unable to get JavaVM instance");
+ }
+ ref_ = env->NewGlobalRef(reader);
+}
+
+gluten::ShuffleStreamReader::~ShuffleStreamReader() {
+ JNIEnv* env = nullptr;
+ attachCurrentThreadAsDaemonOrThrow(vm_, &env);
+ env->DeleteGlobalRef(ref_);
+}
+
+std::shared_ptr<arrow::io::InputStream>
gluten::ShuffleStreamReader::readNextStream(arrow::MemoryPool* pool) {
+ JNIEnv* env = nullptr;
+ attachCurrentThreadAsDaemonOrThrow(vm_, &env);
+
+ jobject jniIn = env->CallObjectMethod(ref_,
getJniCommonState()->shuffleStreamReaderNextStream());
+ checkException(env);
+ if (jniIn == nullptr) {
+ return nullptr; // No more streams to read
+ }
+ return std::make_shared<JavaInputStreamAdaptor>(env, pool, jniIn);
+}
+
std::unique_ptr<gluten::JniColumnarBatchIterator>
gluten::makeJniColumnarBatchIterator(JNIEnv* env, jobject jColumnarBatchItr,
gluten::Runtime* runtime) {
return std::make_unique<JniColumnarBatchIterator>(env, jColumnarBatchItr,
runtime);
diff --git a/cpp/core/jni/JniCommon.h b/cpp/core/jni/JniCommon.h
index bf5a6a746b..f1113a7f82 100644
--- a/cpp/core/jni/JniCommon.h
+++ b/cpp/core/jni/JniCommon.h
@@ -22,9 +22,12 @@
#include <execinfo.h>
#include <jni.h>
+#include <functional>
+
#include "compute/ProtobufUtils.h"
#include "compute/Runtime.h"
#include "memory/AllocationListener.h"
+#include "shuffle/ShuffleReader.h"
#include "shuffle/rss/RssClient.h"
#include "threads/ThreadInitializer.h"
#include "utils/Compression.h"
@@ -151,6 +154,18 @@ static T* jniCastOrThrow(jlong handle) {
}
namespace gluten {
+class ShuffleStreamReader final : public StreamReader {
+ public:
+ ShuffleStreamReader(JNIEnv* env, jobject reader);
+ ~ShuffleStreamReader() override;
+
+ std::shared_ptr<arrow::io::InputStream> readNextStream(arrow::MemoryPool*
pool) override;
+
+ private:
+ JavaVM* vm_{nullptr};
+ jobject ref_{nullptr};
+};
+
class JniCommonState {
public:
virtual ~JniCommonState() = default;
@@ -163,6 +178,14 @@ class JniCommonState {
jmethodID runtimeAwareCtxHandle();
+ jmethodID jniByteInputStreamRead();
+
+ jmethodID jniByteInputStreamTell();
+
+ jmethodID jniByteInputStreamClose();
+
+ jmethodID shuffleStreamReaderNextStream();
+
JavaVM* getJavaVM() const {
return vm_;
}
@@ -173,6 +196,14 @@ class JniCommonState {
jclass runtimeAwareClass_;
jmethodID runtimeAwareCtxHandle_;
+ jclass jniByteInputStreamClass_;
+ jmethodID jniByteInputStreamRead_;
+ jmethodID jniByteInputStreamTell_;
+ jmethodID jniByteInputStreamClose_;
+
+ jclass shuffleStreamReaderClass_;
+ jmethodID shuffleStreamReaderNextStream_;
+
JavaVM* vm_;
bool initialized_{false};
bool closed_{false};
@@ -186,6 +217,14 @@ inline JniCommonState* getJniCommonState() {
Runtime* getRuntime(JNIEnv* env, jobject runtimeAware);
+using JniInputIteratorFactory = std::function<
+ std::unique_ptr<ColumnarBatchIterator>(JNIEnv* env, jobject iterator,
Runtime* runtime, int32_t iteratorIndex)>;
+
+void registerJniInputIteratorFactory(const std::string& runtimeKind,
JniInputIteratorFactory factory);
+
+std::unique_ptr<ColumnarBatchIterator>
+createJniInputIterator(JNIEnv* env, jobject iterator, Runtime* runtime,
int32_t iteratorIndex);
+
// Safe version of JNI {Get|Release}<PrimitiveType>ArrayElements routines.
// SafeNativeArray would release the managed array elements automatically
// during destruction.
diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
index e4d674bd74..91b35c822e 100644
--- a/cpp/core/jni/JniWrapper.cc
+++ b/cpp/core/jni/JniWrapper.cc
@@ -27,7 +27,6 @@
#include "shuffle/ShuffleReader.h"
#include "shuffle/ShuffleWriter.h"
#include "shuffle/Utils.h"
-#include "utils/ArrowStatus.h"
#include "utils/StringUtil.h"
#include <arrow/c/bridge.h>
@@ -51,11 +50,6 @@ jmethodID jniUnsafeByteBufferAllocate;
jmethodID jniUnsafeByteBufferAddress;
jmethodID jniUnsafeByteBufferSize;
-jclass jniByteInputStreamClass;
-jmethodID jniByteInputStreamRead;
-jmethodID jniByteInputStreamTell;
-jmethodID jniByteInputStreamClose;
-
jclass splitResultClass;
jmethodID splitResultConstructor;
@@ -68,9 +62,6 @@ jclass shuffleReaderMetricsClass;
jmethodID shuffleReaderMetricsSetDecompressTime;
jmethodID shuffleReaderMetricsSetDeserializeTime;
-jclass shuffleStreamReaderClass;
-jmethodID shuffleStreamReaderNextStream;
-
jbyteArray toJByteArray(JNIEnv* env, const std::vector<uint8_t>& bytes, const
std::string& context) {
GLUTEN_CHECK(
bytes.size() <= static_cast<size_t>(std::numeric_limits<jsize>::max()),
@@ -82,82 +73,6 @@ jbyteArray toJByteArray(JNIEnv* env, const
std::vector<uint8_t>& bytes, const st
return out;
}
-class JavaInputStreamAdaptor final : public arrow::io::InputStream {
- public:
- JavaInputStreamAdaptor(JNIEnv* env, arrow::MemoryPool* pool, jobject jniIn)
: pool_(pool) {
- // IMPORTANT: DO NOT USE LOCAL REF IN DIFFERENT THREAD
- if (env->GetJavaVM(&vm_) != JNI_OK) {
- std::string errorMessage = "Unable to get JavaVM instance";
- throw GlutenException(errorMessage);
- }
- jniIn_ = env->NewGlobalRef(jniIn);
- }
-
- ~JavaInputStreamAdaptor() override {
- try {
- auto status = JavaInputStreamAdaptor::Close();
- if (!status.ok()) {
- LOG(WARNING) << __func__ << " call JavaInputStreamAdaptor::Close()
failed, status:" << status.ToString();
- }
- } catch (std::exception& e) {
- LOG(WARNING) << __func__ << " call JavaInputStreamAdaptor::Close() got
exception:" << e.what();
- }
- }
-
- // not thread safe
- arrow::Status Close() override {
- if (closed_) {
- return arrow::Status::OK();
- }
- JNIEnv* env;
- attachCurrentThreadAsDaemonOrThrow(vm_, &env);
- env->CallVoidMethod(jniIn_, jniByteInputStreamClose);
- checkException(env);
- env->DeleteGlobalRef(jniIn_);
- // Do NOT call DetachCurrentThread() here.
- // libhdfs.so caches JNIEnv* in thread-local storage after
AttachCurrentThread.
- // If we detach, libhdfs's TLS cache becomes stale — the next HDFS call via
- // libhdfs returns the stale env, causing SIGSEGV in jni_NewStringUTF.
- // Daemon-attached threads are safe to leave attached; they won't block
JVM shutdown.
- closed_ = true;
- return arrow::Status::OK();
- }
-
- arrow::Result<int64_t> Tell() const override {
- JNIEnv* env;
- attachCurrentThreadAsDaemonOrThrow(vm_, &env);
- jlong told = env->CallLongMethod(jniIn_, jniByteInputStreamTell);
- checkException(env);
- return told;
- }
-
- bool closed() const override {
- return closed_;
- }
-
- arrow::Result<int64_t> Read(int64_t nbytes, void* out) override {
- JNIEnv* env;
- attachCurrentThreadAsDaemonOrThrow(vm_, &env);
- jlong read = env->CallLongMethod(jniIn_, jniByteInputStreamRead,
reinterpret_cast<jlong>(out), nbytes);
- checkException(env);
- return read;
- }
-
- arrow::Result<std::shared_ptr<arrow::Buffer>> Read(int64_t nbytes) override {
- GLUTEN_ASSIGN_OR_THROW(auto buffer, arrow::AllocateResizableBuffer(nbytes,
pool_))
- GLUTEN_ASSIGN_OR_THROW(int64_t bytes_read, Read(nbytes,
buffer->mutable_data()))
- GLUTEN_THROW_NOT_OK(buffer->Resize(bytes_read, false));
- buffer->ZeroPadding();
- return std::move(buffer);
- }
-
- private:
- arrow::MemoryPool* pool_;
- JavaVM* vm_;
- jobject jniIn_;
- bool closed_ = false;
-};
-
/// Internal backend consists of empty implementations of Runtime API and
MemoryManager API.
/// The backend is used for saving contextual objects only.
///
@@ -238,39 +153,6 @@ void internalRuntimeReleaser(Runtime* runtime) {
delete runtime;
}
-class ShuffleStreamReader : public StreamReader {
- public:
- ShuffleStreamReader(JNIEnv* env, jobject reader) {
- if (env->GetJavaVM(&vm_) != JNI_OK) {
- throw GlutenException("Unable to get JavaVM instance");
- }
- ref_ = env->NewGlobalRef(reader);
- }
-
- ~ShuffleStreamReader() override {
- JNIEnv* env = nullptr;
- attachCurrentThreadAsDaemonOrThrow(vm_, &env);
- env->DeleteGlobalRef(ref_);
- }
-
- std::shared_ptr<arrow::io::InputStream> readNextStream(arrow::MemoryPool*
pool) override {
- JNIEnv* env = nullptr;
- attachCurrentThreadAsDaemonOrThrow(vm_, &env);
-
- jobject jniIn = env->CallObjectMethod(ref_, shuffleStreamReaderNextStream);
- checkException(env);
- if (jniIn == nullptr) {
- return nullptr; // No more streams to read
- }
- std::shared_ptr<arrow::io::InputStream> in =
std::make_shared<JavaInputStreamAdaptor>(env, pool, jniIn);
- return in;
- }
-
- private:
- JavaVM* vm_;
- jobject ref_;
-};
-
} // namespace
#ifdef __cplusplus
@@ -298,11 +180,6 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) {
jniUnsafeByteBufferAddress = env->GetMethodID(jniUnsafeByteBufferClass,
"address", "()J");
jniUnsafeByteBufferSize = env->GetMethodID(jniUnsafeByteBufferClass, "size",
"()J");
- jniByteInputStreamClass = createGlobalClassReferenceOrError(env,
"Lorg/apache/gluten/vectorized/JniByteInputStream;");
- jniByteInputStreamRead = getMethodIdOrError(env, jniByteInputStreamClass,
"read", "(JJ)J");
- jniByteInputStreamTell = getMethodIdOrError(env, jniByteInputStreamClass,
"tell", "()J");
- jniByteInputStreamClose = getMethodIdOrError(env, jniByteInputStreamClass,
"close", "()V");
-
splitResultClass = createGlobalClassReferenceOrError(env,
"Lorg/apache/gluten/vectorized/GlutenSplitResult;");
splitResultConstructor = getMethodIdOrError(env, splitResultClass, "<init>",
"(JJJJJJJJJJDJ[J[J[J)V");
@@ -322,18 +199,12 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) {
shuffleReaderMetricsSetDeserializeTime =
getMethodIdOrError(env, shuffleReaderMetricsClass, "setDeserializeTime",
"(J)V");
- shuffleStreamReaderClass =
- createGlobalClassReferenceOrError(env,
"Lorg/apache/gluten/vectorized/ShuffleStreamReader;");
- shuffleStreamReaderNextStream = getMethodIdOrError(
- env, shuffleStreamReaderClass, "nextStream",
"()Lorg/apache/gluten/vectorized/JniByteInputStream;");
-
return jniVersion;
}
void JNI_OnUnload(JavaVM* vm, void* reserved) {
JNIEnv* env;
vm->GetEnv(reinterpret_cast<void**>(&env), jniVersion);
- env->DeleteGlobalRef(jniByteInputStreamClass);
env->DeleteGlobalRef(splitResultClass);
env->DeleteGlobalRef(nativeColumnarToRowInfoClass);
env->DeleteGlobalRef(byteArrayClass);
@@ -566,7 +437,7 @@
Java_org_apache_gluten_vectorized_PlanEvaluatorJniWrapper_nativeCreateKernelWith
inputIters.reserve(itersLen);
for (int idx = 0; idx < itersLen; idx++) {
jobject iter = env->GetObjectArrayElement(batchItrArray, idx);
- auto arrayIter = std::make_unique<JniColumnarBatchIterator>(env, iter,
ctx, idx);
+ auto arrayIter = createJniInputIterator(env, iter, ctx, idx);
auto resultIter = std::make_shared<ResultIterator>(std::move(arrayIter));
inputIters.push_back(std::move(resultIter));
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
