run schema fetching in same transaction as arrow data conversion and only once

Author: evertlammerts

src/duckdb_py/arrow/arrow_array_stream.cpp

@@ -14,11 +14,11 @@
 
 namespace duckdb {
 
-void TransformDuckToArrowChunk(ArrowSchema &arrow_schema, ArrowArray &data, py::list &batches) {
+void TransformDuckToArrowChunk(py::object pyarrow_schema, ArrowArray &data, py::list &batches) {
 	py::gil_assert();
 	auto pyarrow_lib_module = py::module::import("pyarrow").attr("lib");
 	auto batch_import_func = pyarrow_lib_module.attr("RecordBatch").attr("_import_from_c");
-	batches.append(batch_import_func(reinterpret_cast<uint64_t>(&data), reinterpret_cast<uint64_t>(&arrow_schema)));
+	batches.append(batch_import_func(reinterpret_cast<uint64_t>(&data), pyarrow_schema));
 }
 
 void VerifyArrowDatasetLoaded() {

src/duckdb_py/arrow/arrow_export_utils.cpp

@@ -17,18 +17,28 @@ namespace duckdb {
 
 namespace pyarrow {
 
-py::object ToArrowTable(const vector<LogicalType> &types, const vector<string> &names, const py::list &batches,
-                        ClientProperties &options) {
+py::object ToPyArrowSchema(ArrowSchema &schema) {
 	py::gil_scoped_acquire acquire;
 
 	auto pyarrow_lib_module = py::module::import("pyarrow").attr("lib");
-	auto from_batches_func = pyarrow_lib_module.attr("Table").attr("from_batches");
 	auto schema_import_func = pyarrow_lib_module.attr("Schema").attr("_import_from_c");
+	return schema_import_func(reinterpret_cast<uint64_t>(&schema));
+}
+
+py::object ToArrowTable(const py::list &batches, py::object pyarrow_schema) {
+	py::gil_scoped_acquire acquire;
+
+	auto pyarrow_lib_module = py::module::import("pyarrow").attr("lib");
+	auto from_batches_func = pyarrow_lib_module.attr("Table").attr("from_batches");
+
+	return py::cast<duckdb::pyarrow::Table>(from_batches_func(batches, pyarrow_schema));
+}
+
+py::object ToArrowTable(const vector<LogicalType> &types, const vector<string> &names, const py::list &batches,
+                        ClientProperties &options) {
 	ArrowSchema schema;
 	ArrowConverter::ToArrowSchema(&schema, types, names, options);
-	auto schema_obj = schema_import_func(reinterpret_cast<uint64_t>(&schema));
-
-	return py::cast<duckdb::pyarrow::Table>(from_batches_func(batches, schema_obj));
+	return ToArrowTable(batches, ToPyArrowSchema(schema));
 }
 
 } // namespace pyarrow

src/duckdb_py/include/duckdb_python/arrow/arrow_array_stream.hpp

@@ -62,7 +62,7 @@ enum class PyArrowObjectType {
 	PolarsLazyFrame
 };
 
-void TransformDuckToArrowChunk(ArrowSchema &arrow_schema, ArrowArray &data, py::list &batches);
+void TransformDuckToArrowChunk(py::object pyarrow_schema, ArrowArray &data, py::list &batches);
 
 PyArrowObjectType GetArrowType(const py::handle &obj);
 

src/duckdb_py/include/duckdb_python/arrow/arrow_export_utils.hpp

@@ -6,9 +6,13 @@ namespace duckdb {
 
 namespace pyarrow {
 
+py::object ToPyArrowSchema(ArrowSchema &schema);
+
 py::object ToArrowTable(const vector<LogicalType> &types, const vector<string> &names, const py::list &batches,
                         ClientProperties &options);
 
+py::object ToArrowTable(const py::list &batches, py::object pyarrow_schema);
+
 } // namespace pyarrow
 
 } // namespace duckdb

src/duckdb_py/pyrelation.cpp

@@ -23,6 +23,59 @@
 #include "duckdb/common/arrow/physical_arrow_collector.hpp"
 #include "duckdb_python/arrow/arrow_export_utils.hpp"
 
+namespace {
+
+// A helper for arrow conversion. We want to be able to fetch a result's schema in the same transaction that
+// creates the result, so we have to wrap both calls in the same transaction. This helper always reverts the
+// transaction if we haven't committed it explicitly. Note that this is not the same as RunFunctionInTransaction:
+// we run _queries_ in a transaction (where each query acquires the context lock) while RFIT runs a function
+// while holding the context lock for that duration.
+// Note: this is a workaround that is intended to be temporary. We should really just cache the schema in the
+// ArrowQueryResult.
+
+void RunOrThrow(duckdb::ClientContext &context, const char *sql) {
+	auto result = context.Query(sql, duckdb::QueryParameters(false));
+	if (result->HasError()) {
+		result->ThrowError();
+	}
+}
+
+class ArrowConversionTransaction {
+public:
+	explicit ArrowConversionTransaction(duckdb::ClientContext &context_p) : context(context_p), owns(false) {
+		auto &txn = context.transaction;
+		if (txn.IsAutoCommit() && !txn.HasActiveTransaction()) {
+			RunOrThrow(context, "BEGIN TRANSACTION");
+			owns = true;
+		}
+	}
+
+	~ArrowConversionTransaction() {
+		if (owns) {
+			try {
+				RunOrThrow(context, "ROLLBACK");
+			} catch (...) { // NOLINT
+			}
+		}
+	}
+
+	void Commit() {
+		if (owns) {
+			RunOrThrow(context, "COMMIT");
+			owns = false;
+		}
+	}
+
+	ArrowConversionTransaction(const ArrowConversionTransaction &) = delete;
+	ArrowConversionTransaction &operator=(const ArrowConversionTransaction &) = delete;
+
+private:
+	duckdb::ClientContext &context;
+	bool owns;
+};
+
+} // namespace
+
 namespace duckdb {
 
 DuckDBPyRelation::DuckDBPyRelation(shared_ptr<Relation> rel_p) : rel(std::move(rel_p)) {
@@ -961,10 +1014,22 @@ PandasDataFrame DuckDBPyRelation::FetchDFChunk(idx_t vectors_per_chunk, bool dat
 }
 
 duckdb::pyarrow::Table DuckDBPyRelation::ToArrowTableInternal(idx_t batch_size, bool to_polars) {
+	if (!result && !rel) {
+		return py::none();
+	}
+	// Make sure we have a valid client context
+	shared_ptr<ClientContext> context;
+	if (rel) {
+		context = rel->context->GetContext();
+	} else if (auto cc = result->GetClientProperties().client_context) {
+		context = cc->shared_from_this();
+	} else {
+		throw ConnectionException("Cannot fetch an arrow table without a valid connection");
+	}
+	// Start (or piggyback on) a transaction for the conversion
+	ArrowConversionTransaction conversion_txn(*context);
+
 	if (!result) {
-		if (!rel) {
-			return py::none();
-		}
 		auto &config = ClientConfig::GetConfig(*rel->context->GetContext());
 		ScopedConfigSetting scoped_setting(
 		    config,
@@ -979,6 +1044,8 @@ duckdb::pyarrow::Table DuckDBPyRelation::ToArrowTableInternal(idx_t batch_size,
 	AssertResultOpen();
 	auto res = result->FetchArrowTable(batch_size, to_polars);
 	result = nullptr;
+	// We must commit the transaction before returning
+	conversion_txn.Commit();
 	return res;
 }
 

src/duckdb_py/pyresult.cpp

@@ -526,21 +526,19 @@ duckdb::pyarrow::Table DuckDBPyResult::FetchArrowTable(idx_t rows_per_batch, boo
 		QueryResult::DeduplicateColumns(names);
 	}
 
+	// Fetch the schema once
+	ArrowSchema arrow_schema;
+	ArrowConverter::ToArrowSchema(&arrow_schema, result->types, names, result->client_properties);
+	auto pyarrow_schema = pyarrow::ToPyArrowSchema(arrow_schema);
+
 	py::list batches;
 	if (result->type == QueryResultType::ARROW_RESULT) {
 		auto &arrow_result = result->Cast<ArrowQueryResult>();
 		auto arrays = arrow_result.ConsumeArrays();
 		for (auto &array : arrays) {
-			ArrowSchema arrow_schema;
-			auto result_names = arrow_result.names;
-			if (to_polars) {
-				QueryResult::DeduplicateColumns(result_names);
-			}
 			ArrowArray data = array->arrow_array;
 			array->arrow_array.release = nullptr;
-			ArrowConverter::ToArrowSchema(&arrow_schema, arrow_result.types, result_names,
-			                              arrow_result.client_properties);
-			TransformDuckToArrowChunk(arrow_schema, data, batches);
+			TransformDuckToArrowChunk(pyarrow_schema, data, batches);
 		}
 	} else {
 		// STREAM_RESULT: pull the live stream directly into Arrow batches.
@@ -558,17 +556,11 @@ duckdb::pyarrow::Table DuckDBPyResult::FetchArrowTable(idx_t rows_per_batch, boo
 			if (count == 0) {
 				break;
 			}
-			ArrowSchema arrow_schema;
-			auto result_names = result->names;
-			if (to_polars) {
-				QueryResult::DeduplicateColumns(result_names);
-			}
-			ArrowConverter::ToArrowSchema(&arrow_schema, result->types, result_names, result->client_properties);
-			TransformDuckToArrowChunk(arrow_schema, data, batches);
+			TransformDuckToArrowChunk(pyarrow_schema, data, batches);
 		}
 	}
 
-	return pyarrow::ToArrowTable(result->types, names, std::move(batches), result->client_properties);
+	return pyarrow::ToArrowTable(std::move(batches), pyarrow_schema);
 }
 
 ArrowArrayStream DuckDBPyResult::FetchArrowArrayStream(idx_t rows_per_batch) {

src/duckdb_py/python_udf.cpp

@@ -25,13 +25,14 @@ static py::list ConvertToSingleBatch(vector<LogicalType> &types, vector<string>
                                      ClientProperties &options, ClientContext &context) {
 	ArrowSchema schema;
 	ArrowConverter::ToArrowSchema(&schema, types, names, options);
+	auto pyarrow_schema = pyarrow::ToPyArrowSchema(schema);
 
 	py::list single_batch;
 	ArrowAppender appender(types, STANDARD_VECTOR_SIZE, options,
 	                       ArrowTypeExtensionData::GetExtensionTypes(context, types));
 	appender.Append(input, 0, input.size(), input.size());
 	auto array = appender.Finalize();
-	TransformDuckToArrowChunk(schema, array, single_batch);
+	TransformDuckToArrowChunk(pyarrow_schema, array, single_batch);
 	return single_batch;
 }