Scalar functions: simplify JNI handling

This PR is a follow-up to #630 and #637.

It removes JNI utilities specific to scalar functions in favour of more
generic `GlobalRefHolder` utility.

Testing: no functional changes, no new tests

Author: staticlibs

CMakeLists.txt

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

CMakeLists.txt.in

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

duckdb_java.def

@@ -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
@@ -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

duckdb_java.exp

@@ -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

duckdb_java.map

@@ -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;
@@ -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;

src/jni/bindings_scalar_function.cpp

@@ -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) {
@@ -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);
@@ -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) {
@@ -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) {
@@ -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);
@@ -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) {
@@ -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);
@@ -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);
@@ -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) {
@@ -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) {

src/jni/holders.cpp

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