Support fetching query results in chunks

CMakeLists.txt

@@ -752,6 +752,7 @@ add_library(duckdb_java SHARED
   src/jni/bindings_common.cpp
   src/jni/bindings_data_chunk.cpp
   src/jni/bindings_logical_type.cpp
+  src/jni/bindings_result.cpp
   src/jni/bindings_scalar_function.cpp
   src/jni/bindings_table_function.cpp
   src/jni/bindings_table_function_bind.cpp

CMakeLists.txt.in

@@ -113,6 +113,7 @@ add_library(duckdb_java SHARED
   src/jni/bindings_common.cpp
   src/jni/bindings_data_chunk.cpp
   src/jni/bindings_logical_type.cpp
+  src/jni/bindings_result.cpp
   src/jni/bindings_scalar_function.cpp
   src/jni/bindings_table_function.cpp
   src/jni/bindings_table_function_bind.cpp

duckdb_java.def

@@ -30,6 +30,7 @@ Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1db_1address
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1extension_1type
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1disconnect
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute
+Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1capi
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1pending
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1fetch
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1cast_1result_1to_1strings
@@ -176,6 +177,13 @@ Java_org_duckdb_DuckDBBindings_duckdb_1get_1timestamp_1ms
 Java_org_duckdb_DuckDBBindings_duckdb_1get_1timestamp_1ns
 Java_org_duckdb_DuckDBBindings_duckdb_1get_1varchar
 
+Java_org_duckdb_DuckDBBindings_duckdb_1destroy_1result
+Java_org_duckdb_DuckDBBindings_duckdb_1fetch_1chunk
+Java_org_duckdb_DuckDBBindings_duckdb_1column_1name
+Java_org_duckdb_DuckDBBindings_duckdb_1column_1type
+Java_org_duckdb_DuckDBBindings_duckdb_1column_1count
+Java_org_duckdb_DuckDBBindings_duckdb_1result_1error
+
 duckdb_adbc_init
 duckdb_add_aggregate_function_to_set
 duckdb_add_replacement_scan

duckdb_java.exp

@@ -27,6 +27,7 @@ _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1db_1address
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1extension_1type
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1disconnect
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute
+_Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1capi
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1pending
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1fetch
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1cast_1result_1to_1strings
@@ -173,6 +174,13 @@ _Java_org_duckdb_DuckDBBindings_duckdb_1get_1timestamp_1ms
 _Java_org_duckdb_DuckDBBindings_duckdb_1get_1timestamp_1ns
 _Java_org_duckdb_DuckDBBindings_duckdb_1get_1varchar
 
+_Java_org_duckdb_DuckDBBindings_duckdb_1destroy_1result
+_Java_org_duckdb_DuckDBBindings_duckdb_1fetch_1chunk
+_Java_org_duckdb_DuckDBBindings_duckdb_1column_1name
+_Java_org_duckdb_DuckDBBindings_duckdb_1column_1type
+_Java_org_duckdb_DuckDBBindings_duckdb_1column_1count
+_Java_org_duckdb_DuckDBBindings_duckdb_1result_1error
+
 _duckdb_adbc_init
 _duckdb_add_aggregate_function_to_set
 _duckdb_add_replacement_scan

duckdb_java.map

@@ -29,6 +29,7 @@ DUCKDB_JAVA {
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1extension_1type;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1disconnect;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute;
+    Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1capi;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1pending;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1fetch;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1cast_1result_1to_1strings;
@@ -175,6 +176,13 @@ DUCKDB_JAVA {
     Java_org_duckdb_DuckDBBindings_duckdb_1get_1timestamp_1ns;
     Java_org_duckdb_DuckDBBindings_duckdb_1get_1varchar;
 
+    Java_org_duckdb_DuckDBBindings_duckdb_1destroy_1result;
+    Java_org_duckdb_DuckDBBindings_duckdb_1fetch_1chunk;
+    Java_org_duckdb_DuckDBBindings_duckdb_1column_1name;
+    Java_org_duckdb_DuckDBBindings_duckdb_1column_1type;
+    Java_org_duckdb_DuckDBBindings_duckdb_1column_1count;
+    Java_org_duckdb_DuckDBBindings_duckdb_1result_1error;
+
     duckdb_adbc_init;
     duckdb_add_aggregate_function_to_set;
     duckdb_add_replacement_scan;

src/jni/bindings_result.cpp

@@ -0,0 +1,144 @@
+#include "bindings.hpp"
+#include "refs.hpp"
+#include "util.hpp"
+
+#include <cstring>
+#include <vector>
+
+duckdb_result *result_buf_to_result(JNIEnv *env, jobject result_buf) {
+
+	if (result_buf == nullptr) {
+		env->ThrowNew(J_SQLException, "Invalid result buffer");
+		return nullptr;
+	}
+
+	duckdb_result *result = reinterpret_cast<duckdb_result *>(env->GetDirectBufferAddress(result_buf));
+	if (result == nullptr) {
+		env->ThrowNew(J_SQLException, "Invalid result");
+		return nullptr;
+	}
+
+	return result;
+}
+
+/*
+ * Class:     org_duckdb_DuckDBBindings
+ * Method:    duckdb_destroy_result
+ * Signature: (Ljava/nio/ByteBuffer;)V
+ */
+JNIEXPORT void JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1destroy_1result(JNIEnv *env, jclass, jobject result) {
+
+	duckdb_result *res = result_buf_to_result(env, result);
+	if (env->ExceptionCheck()) {
+		return;
+	}
+
+	duckdb_result *ptr = res;
+	duckdb_destroy_result(res);
+	delete ptr;
+}
+
+/*
+ * Class:     org_duckdb_DuckDBBindings
+ * Method:    duckdb_fetch_chunk
+ * Signature: (Ljava/nio/ByteBuffer;)Ljava/nio/ByteBuffer;
+ */
+JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1fetch_1chunk(JNIEnv *env, jclass, jobject result) {
+
+	duckdb_result *res = result_buf_to_result(env, result);
+	if (env->ExceptionCheck()) {
+		return nullptr;
+	}
+
+	duckdb_data_chunk chunk = duckdb_fetch_chunk(*res);
+
+	return make_ptr_buf(env, chunk);
+}
+
+/*
+ * Class:     org_duckdb_DuckDBBindings
+ * Method:    duckdb_column_name
+ * Signature: (Ljava/nio/ByteBuffer;J)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1column_1name(JNIEnv *env, jclass, jobject result,
+                                                                                 jlong col) {
+
+	duckdb_result *res = result_buf_to_result(env, result);
+	if (env->ExceptionCheck()) {
+		return nullptr;
+	}
+	idx_t col_idx = jlong_to_idx(env, col);
+	if (env->ExceptionCheck()) {
+		return nullptr;
+	}
+
+	const char *name = duckdb_column_name(res, col_idx);
+	if (name == nullptr) {
+		return nullptr;
+	}
+
+	idx_t len = static_cast<idx_t>(std::strlen(name));
+
+	return make_jbyteArray(env, name, len);
+}
+
+/*
+ * Class:     org_duckdb_DuckDBBindings
+ * Method:    duckdb_column_type
+ * Signature: (Ljava/nio/ByteBuffer;J)I
+ */
+JNIEXPORT jint JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1column_1type(JNIEnv *env, jclass, jobject result,
+                                                                           jlong col) {
+
+	duckdb_result *res = result_buf_to_result(env, result);
+	if (env->ExceptionCheck()) {
+		return -1;
+	}
+	idx_t col_idx = jlong_to_idx(env, col);
+	if (env->ExceptionCheck()) {
+		return -1;
+	}
+
+	duckdb_type dt = duckdb_column_type(res, col_idx);
+
+	return static_cast<jint>(dt);
+}
+
+/*
+ * Class:     org_duckdb_DuckDBBindings
+ * Method:    duckdb_column_count
+ * Signature: (Ljava/nio/ByteBuffer;)J
+ */
+JNIEXPORT jlong JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1column_1count(JNIEnv *env, jclass, jobject result) {
+
+	duckdb_result *res = result_buf_to_result(env, result);
+	if (env->ExceptionCheck()) {
+		return -1;
+	}
+
+	idx_t count = duckdb_column_count(res);
+
+	return static_cast<jlong>(count);
+}
+
+/*
+ * Class:     org_duckdb_DuckDBBindings
+ * Method:    duckdb_result_error
+ * Signature: (Ljava/nio/ByteBuffer;)[B
+ */
+JNIEXPORT jbyteArray JNICALL Java_org_duckdb_DuckDBBindings_duckdb_1result_1error(JNIEnv *env, jclass, jobject result) {
+
+	duckdb_result *res = result_buf_to_result(env, result);
+	if (env->ExceptionCheck()) {
+		return nullptr;
+	}
+
+	const char *error_msg = duckdb_result_error(res);
+	if (error_msg == nullptr) {
+		return nullptr;
+	}
+
+	idx_t len = static_cast<idx_t>(std::strlen(error_msg));
+
+	return make_jbyteArray(env, error_msg, len);
+}

src/jni/duckdb_java.cpp

@@ -13,6 +13,7 @@ extern "C" {
 #include "duckdb/function/scalar/variant_utils.hpp"
 #include "duckdb/function/table/arrow.hpp"
 #include "duckdb/main/appender.hpp"
+#include "duckdb/main/capi/capi_internal.hpp"
 #include "duckdb/main/client_context.hpp"
 #include "duckdb/main/client_data.hpp"
 #include "duckdb/main/database_manager.hpp"
@@ -268,13 +269,13 @@ jobject _duckdb_jdbc_pending_query(JNIEnv *env, jclass, jobject conn_ref_buf, jb
 	return env->NewDirectByteBuffer(pending_ref.release(), 0);
 }
 
-jobject _duckdb_jdbc_execute(JNIEnv *env, jclass, jobject stmt_ref_buf, jobjectArray params) {
+static duckdb::unique_ptr<QueryResult> execute_prepared_statement(JNIEnv *env, jobject stmt_ref_buf,
+                                                                  jobjectArray params, bool stream_results) {
 	auto stmt_ref = reinterpret_cast<StatementHolder *>(env->GetDirectBufferAddress(stmt_ref_buf));
 	if (!stmt_ref) {
 		throw InvalidInputException("Invalid statement");
 	}
 
-	auto res_ref = make_uniq<ResultHolder>();
 	duckdb::vector<Value> duckdb_params;
 
 	idx_t param_len = env->GetArrayLength(params);
@@ -293,20 +294,38 @@ jobject _duckdb_jdbc_execute(JNIEnv *env, jclass, jobject stmt_ref_buf, jobjectA
 		}
 	}
 
+	auto res = stmt_ref->stmt->Execute(duckdb_params, stream_results);
+	if (res->HasError()) {
+		std::string error_msg = std::string(res->GetError());
+		duckdb::ExceptionType error_type = res->GetErrorType();
+		jclass exc_type = duckdb::ExceptionType::INTERRUPT == error_type ? J_SQLTimeoutException : J_SQLException;
+		env->ThrowNew(exc_type, error_msg.c_str());
+		return nullptr;
+	}
+	return res;
+}
+
+jobject _duckdb_jdbc_execute(JNIEnv *env, jclass, jobject stmt_ref_buf, jobjectArray params) {
+	auto stmt_ref = reinterpret_cast<StatementHolder *>(env->GetDirectBufferAddress(stmt_ref_buf));
+	if (!stmt_ref) {
+		throw InvalidInputException("Invalid statement");
+	}
 	Value result;
 	bool stream_results =
 	    stmt_ref->stmt->context->TryGetCurrentSetting("jdbc_stream_results", result) ? result.GetValue<bool>() : false;
+	auto res_ref = make_uniq<ResultHolder>();
+	res_ref->res = execute_prepared_statement(env, stmt_ref_buf, params, stream_results);
+	return env->NewDirectByteBuffer(res_ref.release(), 0);
+}
 
-	res_ref->res = stmt_ref->stmt->Execute(duckdb_params, stream_results);
-	if (res_ref->res->HasError()) {
-		std::string error_msg = std::string(res_ref->res->GetError());
-		duckdb::ExceptionType error_type = res_ref->res->GetErrorType();
-		res_ref->res = nullptr;
-		jclass exc_type = duckdb::ExceptionType::INTERRUPT == error_type ? J_SQLTimeoutException : J_SQLException;
-		env->ThrowNew(exc_type, error_msg.c_str());
+jobject _duckdb_jdbc_execute_capi(JNIEnv *env, jclass, jobject stmt_ref_buf, jobjectArray params) {
+	auto res_ptr = execute_prepared_statement(env, stmt_ref_buf, params, true);
+	if (!res_ptr) {
 		return nullptr;
 	}
-	return env->NewDirectByteBuffer(res_ref.release(), 0);
+	auto out = make_uniq<duckdb_result>();
+	DuckDBTranslateResult(std::move(res_ptr), out.get());
+	return env->NewDirectByteBuffer(out.release(), 0);
 }
 
 jobject _duckdb_jdbc_execute_pending(JNIEnv *env, jclass, jobject pending_ref_buf) {

src/jni/functions.cpp

@@ -206,6 +206,17 @@ JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute(JNI
 	}
 }
 
+JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1capi(JNIEnv * env, jclass param0, jobject param1, jobjectArray param2) {
+	try {
+		return _duckdb_jdbc_execute_capi(env, param0, param1, param2);
+	} catch (const std::exception &e) {
+		duckdb::ErrorData error(e);
+		ThrowJNI(env, error.Message().c_str());
+
+		return nullptr;
+	}
+}
+
 JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1pending(JNIEnv * env, jclass param0, jobject param1) {
 	try {
 		return _duckdb_jdbc_execute_pending(env, param0, param1);

src/jni/functions.hpp

@@ -85,6 +85,10 @@ jobject _duckdb_jdbc_execute(JNIEnv * env, jclass param0, jobject param1, jobjec
 
 JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute(JNIEnv * env, jclass param0, jobject param1, jobjectArray param2);
 
+jobject _duckdb_jdbc_execute_capi(JNIEnv * env, jclass param0, jobject param1, jobjectArray param2);
+
+JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1capi(JNIEnv * env, jclass param0, jobject param1, jobjectArray param2);
+
 jobject _duckdb_jdbc_execute_pending(JNIEnv * env, jclass param0, jobject param1);
 
 JNIEXPORT jobject JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1execute_1pending(JNIEnv * env, jclass param0, jobject param1);

src/main/java/org/duckdb/DuckDBBindings.java

@@ -279,6 +279,20 @@ static native int duckdb_appender_create_ext(ByteBuffer connection, byte[] catal
 
     static native byte[] duckdb_get_varchar(ByteBuffer value);
 
+    // result
+
+    static native void duckdb_destroy_result(ByteBuffer result);
+
+    static native ByteBuffer duckdb_fetch_chunk(ByteBuffer result);
+
+    static native byte[] duckdb_column_name(ByteBuffer result, long col);
+
+    static native int duckdb_column_type(ByteBuffer result, long col);
+
+    static native long duckdb_column_count(ByteBuffer result);
+
+    static native byte[] duckdb_result_error(ByteBuffer result);
+
     enum CAPIType {
         DUCKDB_TYPE_INVALID(0, 0),
         // bool