Restore Arrow streams support

.github/composite-actions/linux-tests-host/action.yml

@@ -4,7 +4,7 @@ description: 'Linux test run in host OS'
 runs:
   using: "composite"
   steps:
-    - name: Setup JDK 26
+    - name: Setup JDK
       uses: actions/setup-java@v6
       with:
         distribution: 'corretto'
@@ -18,3 +18,8 @@ runs:
         java                                                \
           -cp ./build/release/duckdb_jdbc_tests.jar:./build/release/duckdb_jdbc.jar \
           org.duckdb.TestDuckDBJDBC
+
+    - name: Arrow Tests Host
+      shell: bash
+      run: |
+        ./src/test/external/test_arrow.sh

.gitignore

@@ -1,3 +1,6 @@
-build
-.idea
-cmake-build-debug
+/build
+/.idea
+/.local
+/.arrow
+*.iml
+test1.db

CMakeLists.txt

@@ -651,6 +651,7 @@ if(JEMALLOC_ENABLED)
 endif()
 
 add_library(duckdb_java SHARED
+  src/jni/arrow.cpp
   src/jni/bindings_appender.cpp
   src/jni/bindings_common.cpp
   src/jni/bindings_data_chunk.cpp

CMakeLists.txt.in

@@ -131,6 +131,7 @@ if(JEMALLOC_ENABLED)
 endif()
 
 add_library(duckdb_java SHARED
+  src/jni/arrow.cpp
   src/jni/bindings_appender.cpp
   src/jni/bindings_common.cpp
   src/jni/bindings_data_chunk.cpp

data/arrow/arrow_jars.txt

@@ -0,0 +1,14 @@
+https://repo1.maven.org/maven2/org/apache/arrow/arrow-c-data/17.0.0/arrow-c-data-17.0.0.jar
+https://repo1.maven.org/maven2/org/apache/arrow/arrow-format/17.0.0/arrow-format-17.0.0.jar
+https://repo1.maven.org/maven2/org/apache/arrow/arrow-memory-core/17.0.0/arrow-memory-core-17.0.0.jar
+https://repo1.maven.org/maven2/org/apache/arrow/arrow-memory-unsafe/17.0.0/arrow-memory-unsafe-17.0.0.jar
+https://repo1.maven.org/maven2/org/apache/arrow/arrow-vector/17.0.0/arrow-vector-17.0.0.jar
+https://repo1.maven.org/maven2/com/fasterxml/jackson/core/jackson-annotations/2.17.1/jackson-annotations-2.17.1.jar
+https://repo1.maven.org/maven2/com/fasterxml/jackson/core/jackson-core/2.17.1/jackson-core-2.17.1.jar
+https://repo1.maven.org/maven2/com/fasterxml/jackson/core/jackson-databind/2.17.1/jackson-databind-2.17.1.jar
+https://repo1.maven.org/maven2/com/fasterxml/jackson/datatype/jackson-datatype-jsr310/2.17.1/jackson-datatype-jsr310-2.17.1.jar
+https://repo1.maven.org/maven2/com/google/code/findbugs/jsr305/3.0.2/jsr305-3.0.2.jar
+https://repo1.maven.org/maven2/com/google/flatbuffers/flatbuffers-java/24.3.25/flatbuffers-java-24.3.25.jar
+https://repo1.maven.org/maven2/commons-codec/commons-codec/1.17.0/commons-codec-1.17.0.jar
+https://repo1.maven.org/maven2/org/slf4j/slf4j-api/2.0.13/slf4j-api-2.0.13.jar
+https://repo1.maven.org/maven2/org/checkerframework/checker-qual/3.44.0/checker-qual-3.44.0.jar

duckdb_java.def

@@ -20,6 +20,9 @@ Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1begin_1row
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1close
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1end_1row
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1flush
+Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1close_1streams
+Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1register
+Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1stream
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1connect
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1appender
 Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1db_1ref

duckdb_java.exp

@@ -17,6 +17,9 @@ _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1begin_1row
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1close
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1end_1row
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1flush
+_Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1close_1streams
+_Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1register
+_Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1stream
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1connect
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1appender
 _Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1db_1ref

duckdb_java.map

@@ -19,6 +19,9 @@ DUCKDB_JAVA {
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1close;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1end_1row;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1flush;
+    Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1close_1streams;
+    Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1register;
+    Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1arrow_1stream;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1connect;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1appender;
     Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1db_1ref;

src/jni/arrow.cpp

@@ -0,0 +1,184 @@
+#include "arrow.hpp"
+
+#include "duckdb/common/arrow/arrow_appender.hpp"
+#include "duckdb/common/arrow/arrow_converter.hpp"
+#include "duckdb/function/table/arrow/arrow_duck_schema.hpp"
+
+using namespace duckdb;
+
+JavaResultArrowStream::JavaResultArrowStream(duckdb::unique_ptr<QueryResult> res_p,
+                                             duckdb::unique_ptr<QueryResultStream<>> stream_p, idx_t batch_size_p)
+    : res(std::move(res_p)), res_stream(std::move(stream_p)), batch_size(batch_size_p) {
+	if (res) {
+		types = res->GetTypes();
+		names = IdentifiersToStrings(res->GetNames());
+		properties = res->client_properties;
+	} else {
+		types = res_stream->GetTypes();
+		names = IdentifiersToStrings(res_stream->GetNames());
+		properties = res_stream->GetClientProperties();
+	}
+	if (properties.client_context) {
+		extension_types = ArrowTypeExtensionData::GetExtensionTypes(*properties.client_context, types);
+	}
+}
+
+void JavaResultArrowStream::Close() {
+	std::lock_guard<std::mutex> guard(lock);
+	closed = true;
+	chunk = nullptr;
+	res = nullptr;
+	res_stream = nullptr;
+}
+
+ArrowArrayStream *JavaResultArrowStream::Export(duckdb::shared_ptr<JavaResultArrowStream> state) {
+	auto exported = new Exported();
+	exported->state = std::move(state);
+	auto &stream = exported->stream;
+	stream.private_data = exported;
+	stream.get_schema = GetSchema;
+	stream.get_next = GetNext;
+	stream.get_last_error = GetLastError;
+	stream.release = Release;
+	return &stream;
+}
+
+JavaResultArrowStream &JavaResultArrowStream::GetState(ArrowArrayStream *stream) {
+	return *reinterpret_cast<Exported *>(stream->private_data)->state;
+}
+
+int JavaResultArrowStream::GetSchema(ArrowArrayStream *stream, ArrowSchema *out) {
+	if (!stream->release) {
+		return -1;
+	}
+	auto &self = GetState(stream);
+	std::lock_guard<std::mutex> guard(self.lock);
+	out->release = nullptr;
+	try {
+		ArrowConverter::ToArrowSchema(out, self.types, self.names, self.properties);
+	} catch (std::exception &e) {
+		self.last_error = ErrorData(e);
+		return -1;
+	}
+	return 0;
+}
+
+int JavaResultArrowStream::GetNext(ArrowArrayStream *stream, ArrowArray *out) {
+	if (!stream->release) {
+		return -1;
+	}
+	auto &self = GetState(stream);
+	std::lock_guard<std::mutex> guard(self.lock);
+	out->release = nullptr;
+	if (self.closed) {
+		self.last_error = ErrorData(ConnectionException("Connection was closed"));
+		return -1;
+	}
+	try {
+		ArrowAppender appender(self.types, self.batch_size, self.properties, self.extension_types);
+		idx_t count = 0;
+		while (count < self.batch_size) {
+			if (!self.chunk || self.chunk_offset >= self.chunk->size()) {
+				if (!self.FetchChunk()) {
+					return -1;
+				}
+				if (!self.chunk) {
+					break;
+				}
+				continue;
+			}
+			auto &chunk = *self.chunk;
+			auto to_append = MinValue(self.batch_size - count, chunk.size() - self.chunk_offset);
+			appender.Append(chunk, self.chunk_offset, self.chunk_offset + to_append, chunk.size());
+			self.chunk_offset += to_append;
+			count += to_append;
+		}
+		if (count > 0) {
+			*out = appender.Finalize();
+		}
+	} catch (std::exception &e) {
+		self.last_error = ErrorData(e);
+		return -1;
+	}
+	return 0;
+}
+
+const char *JavaResultArrowStream::GetLastError(ArrowArrayStream *stream) {
+	if (!stream->release) {
+		return "stream was released";
+	}
+	auto &self = GetState(stream);
+	std::lock_guard<std::mutex> guard(self.lock);
+	return self.last_error.HasError() ? self.last_error.Message().c_str() : nullptr;
+}
+
+void JavaResultArrowStream::Release(ArrowArrayStream *stream) {
+	if (!stream || !stream->release) {
+		return;
+	}
+	stream->release = nullptr;
+	delete reinterpret_cast<Exported *>(stream->private_data);
+}
+
+bool JavaResultArrowStream::FetchChunk() {
+	chunk_offset = 0;
+	if (finished) {
+		chunk = nullptr;
+		return true;
+	}
+	if (res) {
+		chunk = res->Fetch();
+		if (res->HasError()) {
+			last_error = res->GetErrorObject();
+			return false;
+		}
+	} else {
+		chunk = res_stream->Fetch();
+		if (res_stream->HasError()) {
+			last_error = res_stream->GetErrorObject();
+			return false;
+		}
+	}
+	if (!chunk) {
+		finished = true;
+	}
+	return true;
+}
+
+JavaRegisteredArrowStream::JavaRegisteredArrowStream(ArrowArrayStream *source) : stream(*source) {
+	source->release = nullptr;
+}
+
+JavaRegisteredArrowStream::~JavaRegisteredArrowStream() {
+	Close();
+}
+
+void JavaRegisteredArrowStream::Close() {
+	std::lock_guard<std::mutex> guard(lock);
+	if (stream.release) {
+		stream.release(&stream);
+		stream.release = nullptr;
+	}
+}
+
+void JavaRegisteredArrowStream::GetSchema(ArrowSchema &schema) {
+	std::lock_guard<std::mutex> guard(lock);
+	if (!stream.release) {
+		throw InvalidInputException("This stream has been released");
+	}
+	if (stream.get_schema(&stream, &schema) != 0) {
+		auto error = stream.get_last_error(&stream);
+		throw InvalidInputException(error ? error : "Unknown error getting Arrow stream schema");
+	}
+}
+
+duckdb::unique_ptr<ArrowArrayStreamWrapper> JavaRegisteredArrowStream::Take() {
+	std::lock_guard<std::mutex> guard(lock);
+	if (!stream.release) {
+		throw InvalidInputException("This stream has been released");
+	}
+	auto res = make_uniq<ArrowArrayStreamWrapper>();
+	res->arrow_array_stream = stream;
+	stream.release = nullptr;
+	return res;
+}

src/jni/arrow.hpp

@@ -0,0 +1,75 @@
+#pragma once
+
+#include "duckdb.hpp"
+#include "duckdb/common/arrow/arrow_wrapper.hpp"
+#include "duckdb/main/query_result_stream.hpp"
+#include "holders.hpp"
+
+#include <mutex>
+
+namespace duckdb {
+class ArrowTypeExtensionData;
+}
+
+//! Exports the chunks of a JDBC result as an Arrow stream of record batches of at most batch_size rows.
+//! The result is moved out of the ResultHolder into this state, that is shared between the exported stream
+//! and the connection. Connection closes the state (releasing the result) if it is closed before the stream.
+class JavaResultArrowStream : public ArrowStreamState {
+public:
+	JavaResultArrowStream(duckdb::unique_ptr<duckdb::QueryResult> res_p,
+	                      duckdb::unique_ptr<duckdb::QueryResultStream<>> stream_p, duckdb::idx_t batch_size_p);
+
+	void Close() override;
+
+	//! Creates a stream that shares the state, the stream is deleted when it is released
+	static ArrowArrayStream *Export(duckdb::shared_ptr<JavaResultArrowStream> state);
+
+private:
+	struct Exported {
+		ArrowArrayStream stream;
+		duckdb::shared_ptr<JavaResultArrowStream> state;
+	};
+
+	static JavaResultArrowStream &GetState(ArrowArrayStream *stream);
+	static int GetSchema(ArrowArrayStream *stream, ArrowSchema *out);
+	static int GetNext(ArrowArrayStream *stream, ArrowArray *out);
+	static const char *GetLastError(ArrowArrayStream *stream);
+	static void Release(ArrowArrayStream *stream);
+
+	//! Loads the next chunk, chunk is null at the end of the result
+	bool FetchChunk();
+
+	std::mutex lock;
+	duckdb::unique_ptr<duckdb::QueryResult> res;
+	duckdb::unique_ptr<duckdb::QueryResultStream<>> res_stream;
+	duckdb::idx_t batch_size;
+	duckdb::vector<duckdb::LogicalType> types;
+	duckdb::vector<std::string> names;
+	duckdb::ClientProperties properties;
+	duckdb::unordered_map<duckdb::idx_t, const duckdb::shared_ptr<duckdb::ArrowTypeExtensionData>> extension_types;
+	duckdb::unique_ptr<duckdb::DataChunk> chunk;
+	duckdb::idx_t chunk_offset = 0;
+	bool finished = false;
+	bool closed = false;
+	duckdb::ErrorData last_error;
+};
+
+//! Arrow stream registered with a connection. The stream is moved out of the Java-owned struct on registration,
+//! it is released when it is consumed by a scan, when the view is dropped or when the connection is closed.
+class JavaRegisteredArrowStream : public ArrowStreamState {
+public:
+	explicit JavaRegisteredArrowStream(ArrowArrayStream *source);
+
+	~JavaRegisteredArrowStream() override;
+
+	void Close() override;
+
+	void GetSchema(ArrowSchema &schema);
+
+	//! A stream can only be consumed by a single scan
+	duckdb::unique_ptr<duckdb::ArrowArrayStreamWrapper> Take();
+
+private:
+	std::mutex lock;
+	ArrowArrayStream stream;
+};