Skip to content

Scalar functions: simplify JNI handling #642

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
staticlibs wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Scalar functions: simplify JNI handling #642

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -597,8 +597,8 @@ add_library(duckdb_java SHARED
src/jni/config.cpp
src/jni/duckdb_java.cpp
src/jni/functions.cpp
src/jni/holders.cpp
src/jni/refs.cpp
src/jni/scalar_functions.cpp
src/jni/types.cpp
src/jni/util.cpp
${DUCKDB_SRC_FILES})
Expand Down
2 changes: 1 addition & 1 deletion CMakeLists.txt.in
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,8 +115,8 @@ add_library(duckdb_java SHARED
src/jni/config.cpp
src/jni/duckdb_java.cpp
src/jni/functions.cpp
src/jni/holders.cpp
src/jni/refs.cpp
src/jni/scalar_functions.cpp
src/jni/types.cpp
src/jni/util.cpp
${DUCKDB_SRC_FILES})
Expand Down
3 changes: 2 additions & 1 deletion duckdb_java.def
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@ Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1return_1type
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1varargs
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1volatile
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1special_1handling
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1extra_1info
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1function
Java_org_duckdb_DuckDBBindings_duckdb_1register_1scalar_1function
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1error
Java_org_duckdb_DuckDBBindings_duckdb_1create_1logical_1type
Expand Down Expand Up @@ -111,7 +113,6 @@ Java_org_duckdb_DuckDBBindings_duckdb_1appender_1column_1count
Java_org_duckdb_DuckDBBindings_duckdb_1appender_1column_1type
Java_org_duckdb_DuckDBBindings_duckdb_1append_1data_1chunk
Java_org_duckdb_DuckDBBindings_duckdb_1append_1default_1to_1chunk
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1function

duckdb_adbc_init
duckdb_add_aggregate_function_to_set
Expand Down
3 changes: 2 additions & 1 deletion duckdb_java.exp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,13 +57,14 @@ _Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1return_1type
_Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1varargs
_Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1volatile
_Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1special_1handling
_Java_org_duckdb_DuckDBBindings_duckdb_1register_1scalar_1function
_Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1error
_Java_org_duckdb_DuckDBBindings_duckdb_1vector_1get_1string
_Java_org_duckdb_DuckDBBindings_duckdb_1vector_1get_1string__Ljava_nio_ByteBuffer_2J
_Java_org_duckdb_DuckDBBindings_duckdb_1create_1logical_1type
_Java_org_duckdb_DuckDBBindings_duckdb_1create_1decimal_1type
_Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1extra_1info
_Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1function
_Java_org_duckdb_DuckDBBindings_duckdb_1register_1scalar_1function
_Java_org_duckdb_DuckDBBindings_duckdb_1get_1type_1id
_Java_org_duckdb_DuckDBBindings_duckdb_1decimal_1width
_Java_org_duckdb_DuckDBBindings_duckdb_1decimal_1scale
Expand Down
3 changes: 2 additions & 1 deletion duckdb_java.map
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ DUCKDB_JAVA {
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1varargs;
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1volatile;
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1special_1handling;
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1extra_1info;
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1function;
Java_org_duckdb_DuckDBBindings_duckdb_1register_1scalar_1function;
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1error;
Java_org_duckdb_DuckDBBindings_duckdb_1vector_1get_1string;
Expand Down Expand Up @@ -110,7 +112,6 @@ DUCKDB_JAVA {
Java_org_duckdb_DuckDBBindings_duckdb_1appender_1column_1type;
Java_org_duckdb_DuckDBBindings_duckdb_1append_1data_1chunk;
Java_org_duckdb_DuckDBBindings_duckdb_1append_1default_1to_1chunk;
Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1function;

duckdb_adbc_init;
duckdb_add_aggregate_function_to_set;
Expand Down
129 changes: 121 additions & 8 deletions src/jni/bindings_scalar_function.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
#include "functions.hpp"
#include "holders.hpp"
#include "refs.hpp"
#include "scalar_functions.hpp"
#include "util.hpp"

static duckdb_scalar_function scalar_function_buf_to_scalar_function(JNIEnv *env, jobject scalar_function_buf) {
Expand Down Expand Up @@ -36,10 +35,20 @@ static duckdb_function_info function_info_buf_to_function_info(JNIEnv *env, jobj
return function_info;
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_create_scalar_function
* Signature: ()Ljava/nio/ByteBuffer;
*/
JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1create_1scalar_1function(JNIEnv *env, jclass) {
return make_ptr_buf(env, duckdb_create_scalar_function());
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_destroy_scalar_function
* Signature: (Ljava/nio/ByteBuffer;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1destroy_1scalar_1function(JNIEnv *env, jclass,
jobject scalar_function) {
auto function = scalar_function_buf_to_scalar_function(env, scalar_function);
Expand All @@ -49,6 +58,11 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1destroy_1scalar_1f
duckdb_destroy_scalar_function(&function);
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_name
* Signature: (Ljava/nio/ByteBuffer;[B)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1name(JNIEnv *env, jclass,
jobject scalar_function,
jbyteArray name) {
Expand All @@ -67,6 +81,11 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1
duckdb_scalar_function_set_name(function, function_name.c_str());
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_add_parameter
* Signature: (Ljava/nio/ByteBuffer;Ljava/nio/ByteBuffer;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1add_1parameter(JNIEnv *env, jclass,
jobject scalar_function,
jobject logical_type) {
Expand All @@ -81,6 +100,11 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1
duckdb_scalar_function_add_parameter(function, type);
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_return_type
* Signature: (Ljava/nio/ByteBuffer;Ljava/nio/ByteBuffer;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1return_1type(
JNIEnv *env, jclass, jobject scalar_function, jobject logical_type) {
auto function = scalar_function_buf_to_scalar_function(env, scalar_function);
Expand All @@ -94,6 +118,11 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1
duckdb_scalar_function_set_return_type(function, type);
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_varargs
* Signature: (Ljava/nio/ByteBuffer;Ljava/nio/ByteBuffer;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1varargs(JNIEnv *env, jclass,
jobject scalar_function,
jobject logical_type) {
Expand All @@ -108,6 +137,11 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1
duckdb_scalar_function_set_varargs(function, type);
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_volatile
* Signature: (Ljava/nio/ByteBuffer;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1volatile(JNIEnv *env, jclass,
jobject scalar_function) {
auto function = scalar_function_buf_to_scalar_function(env, scalar_function);
Expand All @@ -117,6 +151,11 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1
duckdb_scalar_function_set_volatile(function);
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_special_handling
* Signature: (Ljava/nio/ByteBuffer;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1special_1handling(
JNIEnv *env, jclass, jobject scalar_function) {
auto function = scalar_function_buf_to_scalar_function(env, scalar_function);
Expand All @@ -126,6 +165,11 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1
duckdb_scalar_function_set_special_handling(function);
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_register_scalar_function
* Signature: (Ljava/nio/ByteBuffer;Ljava/nio/ByteBuffer;)I
*/
JNIEXPORT jint JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1register_1scalar_1function(JNIEnv *env, jclass,
jobject connection,
jobject scalar_function) {
Expand All @@ -140,16 +184,85 @@ JNIEXPORT jint JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1register_1scalar_1
return static_cast<jint>(duckdb_register_scalar_function(conn, function));
}

JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1function(
JNIEnv *env, jclass, jobject scalar_function_buf, jobject function_j) {
try {
scalar_function_set_function(env, scalar_function_buf, function_j);
} catch (const std::exception &e) {
duckdb::ErrorData error(e);
ThrowJNI(env, error.Message().c_str());
/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_extra_info
* Signature: (Ljava/nio/ByteBuffer;Ljava/lang/Object;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1extra_1info(
JNIEnv *env, jclass, jobject scalar_function, jobject callback) {

if (callback == nullptr) {
env->ThrowNew(J_SQLException, "Specified callback must be not null");
return;
}

auto sf = scalar_function_buf_to_scalar_function(env, scalar_function);
if (env->ExceptionCheck()) {
return;
}

auto callback_holder = std::unique_ptr<GlobalRefHolder>(new GlobalRefHolder(env, callback));
if (callback_holder->vm == nullptr) {
env->ThrowNew(J_SQLException, "Unable to create a global reference to the specified scalar function callback");
return;
}

duckdb_scalar_function_set_extra_info(sf, callback_holder.release(), GlobalRefHolder::destroy);
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_function
* Signature: (Ljava/nio/ByteBuffer;)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1function(JNIEnv *env, jclass,
jobject scalar_function) {
auto sf = scalar_function_buf_to_scalar_function(env, scalar_function);
if (env->ExceptionCheck()) {
return;
}

duckdb_scalar_function_set_function(
sf, [](duckdb_function_info info, duckdb_data_chunk input, duckdb_vector output) {
auto callback_holder = reinterpret_cast<GlobalRefHolder *>(duckdb_scalar_function_get_extra_info(info));
AttachedJNIEnv attached = callback_holder->attach_current_thread();
if (attached.env == nullptr) {
duckdb_scalar_function_set_error(info, "Unable to attach JNI environment");
return;
}
jobject info_buf = make_ptr_buf(attached.env, info);
if (info_buf == nullptr) {
duckdb_scalar_function_set_error(info, "Unable to create function info buffer");
return;
}
LocalRefHolder info_buf_holder(attached.env, info_buf);
jobject input_buf = make_ptr_buf(attached.env, input);
if (input_buf == nullptr) {
duckdb_scalar_function_set_error(info, "Unable to create input buffer");
return;
}
LocalRefHolder input_buf_holder(attached.env, input_buf);
jobject output_buf = make_ptr_buf(attached.env, output);
if (output_buf == nullptr) {
duckdb_scalar_function_set_error(info, "Unable to create output buffer");
return;
}
LocalRefHolder output_buf_holder(attached.env, output_buf);

attached.env->CallVoidMethod(callback_holder->global_ref, J_DuckDBScalarFunctionWrapper_execute, info_buf,
input_buf, output_buf);
if (attached.env->ExceptionCheck()) {
duckdb_scalar_function_set_error(info, "Java callback system error");
}
});
}

/*
* Class: org_duckdb_DuckDBBindings
* Method: duckdb_scalar_function_set_error
* Signature: (Ljava/nio/ByteBuffer;[B)V
*/
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1scalar_1function_1set_1error(JNIEnv *env, jclass,
jobject function_info_buf,
jbyteArray error) {
Expand Down
120 changes: 120 additions & 0 deletions src/jni/holders.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
#include "holders.hpp"

ConnectionHolder *get_connection_ref(JNIEnv *env, jobject conn_ref_buf) {
if (!conn_ref_buf) {
throw duckdb::ConnectionException("Invalid connection buffer ref");
}
auto conn_holder = reinterpret_cast<ConnectionHolder *>(env->GetDirectBufferAddress(conn_ref_buf));
if (!conn_holder) {
throw duckdb::ConnectionException("Invalid connection buffer");
}
return conn_holder;
}

/**
* Throws a SQLException and returns nullptr if a valid Connection can't be retrieved from the buffer.
*/
duckdb::Connection *get_connection(JNIEnv *env, jobject conn_ref_buf) {
auto conn_holder = get_connection_ref(env, conn_ref_buf);
auto conn_ref = conn_holder->connection.get();
if (!conn_ref || !conn_ref->context) {
throw duckdb::ConnectionException("Invalid connection");
}

return conn_ref;
}

duckdb_connection conn_ref_buf_to_conn(JNIEnv *env, jobject conn_ref_buf) {
if (conn_ref_buf == nullptr) {
env->ThrowNew(J_SQLException, "Invalid connection buffer");
return nullptr;
}
auto conn_holder = reinterpret_cast<ConnectionHolder *>(env->GetDirectBufferAddress(conn_ref_buf));
if (conn_holder == nullptr) {
env->ThrowNew(J_SQLException, "Invalid connection holder");
return nullptr;
}
auto conn_ref = conn_holder->connection.get();
if (conn_ref == nullptr || conn_ref->context == nullptr) {
env->ThrowNew(J_SQLException, "Invalid connection");
return nullptr;
}

return reinterpret_cast<duckdb_connection>(conn_ref);
}

AttachedJNIEnv::AttachedJNIEnv() {
}

AttachedJNIEnv::AttachedJNIEnv(JavaVM *vm_in, JNIEnv *env_in, bool need_to_detach_in)
: vm(vm_in), env(env_in), need_to_detach(need_to_detach_in) {
}

AttachedJNIEnv::~AttachedJNIEnv() noexcept {
if (vm == nullptr) {
return;
}
if (need_to_detach) {
vm->DetachCurrentThread();
}
}

GlobalRefHolder::GlobalRefHolder(JNIEnv *env, jobject local_ref) {
if (env->GetJavaVM(&this->vm) != JNI_OK || this->vm == nullptr) {
this->vm = nullptr;
return;
}
if (local_ref != nullptr) {
this->global_ref = env->NewGlobalRef(local_ref);
if (this->global_ref == nullptr) {
this->vm = nullptr;
}
}
}

GlobalRefHolder::~GlobalRefHolder() noexcept {
if (global_ref == nullptr) {
return;
}
AttachedJNIEnv attached = attach_current_thread();
if (attached.env == nullptr) {
return;
}
attached.env->DeleteGlobalRef(global_ref);
}

AttachedJNIEnv GlobalRefHolder::attach_current_thread() {
if (vm == nullptr) {
return AttachedJNIEnv();
}
JNIEnv *env = nullptr;
auto env_status = vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_6);
if (env_status != JNI_OK && env_status != JNI_EDETACHED) {
return AttachedJNIEnv();
}
bool need_to_detach = false;
if (env_status == JNI_EDETACHED) {
auto attach_status = vm->AttachCurrentThread(reinterpret_cast<void **>(&env), nullptr);
if (attach_status != JNI_OK || env == nullptr) {
return AttachedJNIEnv();
}
need_to_detach = true;
}

return AttachedJNIEnv(vm, env, need_to_detach);
}

void GlobalRefHolder::destroy(void *holder_in) noexcept {
auto holder = reinterpret_cast<GlobalRefHolder *>(holder_in);
delete holder;
}

LocalRefHolder::LocalRefHolder(JNIEnv *env_in, jobject local_ref_in) : env(env_in), local_ref(local_ref_in) {
}

LocalRefHolder::~LocalRefHolder() noexcept {
if (env == nullptr || local_ref == nullptr) {
return;
}
env->DeleteLocalRef(local_ref);
}
Loading
Loading