keep connection alive for arrow streams

Author: evertlammerts

external/duckdb

@@ -46,7 +46,7 @@ struct DuckDBPyResult {
 
 	ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch = 1000000);
 	duckdb::pyarrow::RecordBatchReader FetchRecordBatchReader(idx_t rows_per_batch = 1000000);
-	py::object FetchArrowCapsule(idx_t rows_per_batch = 1000000);
+	py::object FetchArrowCapsule(shared_ptr<ClientContext> context, idx_t rows_per_batch = 1000000);
 
 	static py::list GetDescription(const vector<string> &names, const vector<LogicalType> &types);
 

src/duckdb_py/include/duckdb_python/pyresult.hpp

@@ -1006,7 +1006,8 @@ py::object DuckDBPyRelation::ToArrowCapsule(const py::object &requested_schema)
 		ExecuteOrThrow();
 	}
 	AssertResultOpen();
-	auto res = result->FetchArrowCapsule();
+	auto context = rel ? rel->context->GetContext() : nullptr;
+	auto res = result->FetchArrowCapsule(std::move(context));
 	result = nullptr;
 	return res;
 }

src/duckdb_py/pyrelation.cpp

@@ -500,7 +500,8 @@ duckdb::pyarrow::RecordBatchReader DuckDBPyResult::FetchRecordBatchReader(idx_t
 // This avoids the double-materialization that happens when using ResultArrowArrayStreamWrapper
 // with an ArrowQueryResult (which throws NotImplementedException from FetchInternal).
 struct ArrowQueryResultStreamWrapper {
-	ArrowQueryResultStreamWrapper(unique_ptr<QueryResult> result_p) : result(std::move(result_p)), index(0) {
+	ArrowQueryResultStreamWrapper(unique_ptr<QueryResult> result_p, shared_ptr<ClientContext> context_p)
+	    : result(std::move(result_p)), context(std::move(context_p)), index(0) {
 		auto &arrow_result = result->Cast<ArrowQueryResult>();
 		arrays = arrow_result.ConsumeArrays();
 		types = result->types;
@@ -562,6 +563,7 @@ struct ArrowQueryResultStreamWrapper {
 
 	ArrowArrayStream stream;
 	unique_ptr<QueryResult> result;
+	shared_ptr<ClientContext> context;
 	vector<unique_ptr<ArrowArrayWrapper>> arrays;
 	vector<LogicalType> types;
 	vector<string> names;
@@ -570,7 +572,9 @@ struct ArrowQueryResultStreamWrapper {
 	string last_error;
 };
 
-// Destructor for capsules that own a heap-allocated ArrowArrayStream (slow path).
+// Destructor for capsules that own a heap-allocated ArrowArrayStream.
+// If PyCapsule_GetContext is set, it points to a shared_ptr<ClientContext> that
+// keeps the connection's ClientContext alive for the lifetime of the capsule.
 static void ArrowArrayStreamPyCapsuleDestructor(PyObject *object) {
 	auto data = PyCapsule_GetPointer(object, "arrow_array_stream");
 	if (!data) {
@@ -581,25 +585,35 @@ static void ArrowArrayStreamPyCapsuleDestructor(PyObject *object) {
 		stream->release(stream);
 	}
 	delete stream;
+	auto ctx = PyCapsule_GetContext(object);
+	if (ctx) {
+		delete reinterpret_cast<shared_ptr<ClientContext> *>(ctx);
+	}
 }
 
-py::object DuckDBPyResult::FetchArrowCapsule(idx_t rows_per_batch) {
+py::object DuckDBPyResult::FetchArrowCapsule(shared_ptr<ClientContext> context, idx_t rows_per_batch) {
 	if (result && result->type == QueryResultType::ARROW_RESULT) {
 		// Fast path: yield pre-built Arrow arrays directly.
 		// The wrapper is heap-allocated; Release() deletes it via private_data.
 		// We heap-allocate a separate ArrowArrayStream for the capsule so that the capsule
 		// holds a stable pointer even after the wrapper is consumed and deleted by a scan.
-		auto wrapper = new ArrowQueryResultStreamWrapper(std::move(result));
+		auto wrapper = new ArrowQueryResultStreamWrapper(std::move(result), std::move(context));
 		auto stream = new ArrowArrayStream();
 		*stream = wrapper->stream;
 		wrapper->stream.release = nullptr;
 		return py::capsule(stream, "arrow_array_stream", ArrowArrayStreamPyCapsuleDestructor);
 	}
-	// Existing slow path for MaterializedQueryResult / StreamQueryResult
+	// Existing slow path for MaterializedQueryResult / StreamQueryResult.
+	// Keep the ClientContext alive via the capsule's context pointer so that
+	// deferred get_schema / get_next calls can still dereference client_context.
 	auto stream_p = FetchArrowArrayStream(rows_per_batch);
 	auto stream = new ArrowArrayStream();
 	*stream = stream_p;
-	return py::capsule(stream, "arrow_array_stream", ArrowArrayStreamPyCapsuleDestructor);
+	auto capsule = py::capsule(stream, "arrow_array_stream", ArrowArrayStreamPyCapsuleDestructor);
+	if (context) {
+		PyCapsule_SetContext(capsule.ptr(), new shared_ptr<ClientContext>(std::move(context)));
+	}
+	return capsule;
 }
 
 py::list DuckDBPyResult::GetDescription(const vector<string> &names, const vector<LogicalType> &types) {