Unify arrow exports across all query result types (#495)

This PR unifies arrow exports across query result types, and makes sure
we always provide the schema from within a transaction.

We are dealing with 3 arrow export types:
- Arrow Table
- Arrow RecordBatch
- Arrow C Stream

... across 3 result types:
- StreamingQueryResult
- ArrowQueryResult
- StreamingQueryResult

The `StreamingQueryResult` paths are now unified. We re-feed the backing
ColumnDataCollection to the engine for parallel conversion into a
`ArrowQueryResult`, and then we delegate to the corresponding
`ArrowQueryResult` path.

The `ArrowQueryResult` paths deal with materialized data already, and we
have no way to plug into the transaction that generated it. The actual
fix for this is to cache the schema when creating the
`ArrowQueryResult`, during `Finalize`. This is a core change that we
will probably apply in v2.0. The workaround is to fetch the schema in a
separate transaction. For all paths, since we are already dealing with
materialized data, we create an arrow table. Then for the streaming
paths we return the corresponding stream types directly from the table.

The `StreamingQueryResult` paths always have access to a valid
transaction context, and can get the arrow schema on demand even when
that requires catalog access.

As a side effect of this PR, consuming an arrow c stream (reading from
`con.sql(q).__arrow_c_stream__()`) is now lazy, i.e. not materialized.
This makes consumption of course slower, but allows streaming much
larger datasets.

The materialized paths are overall a little faster, and the non-c stream
streaming paths as well.

```
  ┌───────────────────────────────────────────────────┬────────────────────┬───────────────────┬───────────────────┐
  │               benchmark expression                │ wall base→now (ms) │ CPU base→now (ms) │ mem base→now (MB) │
  ├───────────────────────────────────────────────────┼────────────────────┼───────────────────┼───────────────────┤
  │ r=con.sql(q); r.execute(); r.to_arrow_table()     │ 159 → 161          │ 259 → 286         │ 847 → 875         │
  ├───────────────────────────────────────────────────┼────────────────────┼───────────────────┼───────────────────┤
  │ r=con.sql(q); r.execute(); r.to_arrow_reader()    │ 161 → 144          │ 255 → 263         │ 896 → 877         │
  ├───────────────────────────────────────────────────┼────────────────────┼───────────────────┼───────────────────┤
  │ r=con.sql(q); r.execute(); r.__arrow_c_stream__() │ 157 → 136          │ 282 → 235         │ 854 → 881         │
  ├───────────────────────────────────────────────────┼────────────────────┼───────────────────┼───────────────────┤
  │ con.sql(q).to_arrow_table()                       │ 52 → 35            │ 267 → 244         │ 855 → 854         │
  ├───────────────────────────────────────────────────┼────────────────────┼───────────────────┼───────────────────┤
  │ con.execute(q).to_arrow_table()                   │ 202 → 174          │ 212 → 193         │ 548 → 554         │
  ├───────────────────────────────────────────────────┼────────────────────┼───────────────────┼───────────────────┤
  │ con.sql(q).to_arrow_reader()                      │ 186 → 175          │ 199 → 187         │ 552 → 552         │
  ├───────────────────────────────────────────────────┼────────────────────┼───────────────────┼───────────────────┤
  │ con.sql(q).__arrow_c_stream__()                   │ 48 → 173           │ 250 → 189         │ 857 → 554         │
  └───────────────────────────────────────────────────┴────────────────────┴───────────────────┴───────────────────┘

```

Author: evertlammerts

.github/workflows/packaging_wheels.yml

@@ -32,7 +32,7 @@ jobs:
       matrix:
         python: [ cp310, cp311, cp312, cp313, cp314 ]
         platform:
-          - { os: windows-2025,     arch: amd64,      cibw_system: win }
+          - { os: windows-2022,     arch: amd64,      cibw_system: win }
           - { os: windows-11-arm,   arch: ARM64,      cibw_system: win } # cibw requires ARM64 to be uppercase
           - { os: ubuntu-24.04,     arch: x86_64,     cibw_system: manylinux }
           - { os: ubuntu-24.04-arm, arch: aarch64,    cibw_system: manylinux }

.github/workflows/targeted_test.yml

@@ -8,7 +8,7 @@ on:
         required: true
         type: choice
         options:
-          - 'windows-2025'
+          - 'windows-2022'
           - 'windows-11-arm'
           - 'ubuntu-24.04'
           - 'ubuntu-24.04-arm'

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(const 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(const 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/include/duckdb_python/pyresult.hpp

@@ -36,15 +36,13 @@ struct DuckDBPyResult {
 
 	PandasDataFrame FetchDF(bool date_as_object);
 
-	duckdb::pyarrow::Table FetchArrowTable(idx_t rows_per_batch, bool to_polars);
-
 	PandasDataFrame FetchDFChunk(const idx_t vectors_per_chunk = 1, bool date_as_object = false);
 
 	py::dict FetchPyTorch();
 
 	py::dict FetchTF();
 
-	ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch = 1000000);
+	duckdb::pyarrow::Table FetchArrowTable(idx_t rows_per_batch, bool to_polars);
 	duckdb::pyarrow::RecordBatchReader FetchRecordBatchReader(idx_t rows_per_batch = 1000000);
 	py::object FetchArrowCapsule(idx_t rows_per_batch = 1000000);
 
@@ -71,6 +69,19 @@ struct DuckDBPyResult {
 	unique_ptr<DataChunk> FetchNextRaw(QueryResult &result);
 	unique_ptr<NumpyResultConversion> InitializeNumpyConversion(bool pandas = false);
 
+	//! Re-feed an already-MATERIALIZED result (a ColumnDataCollection, e.g. from
+	//! rel.execute()) back through the engine on the user's own context. The eager
+	//! variant installs a PhysicalArrowCollector to produce an ArrowQueryResult
+	//! (parallel); the stream variant produces a lazy StreamQueryResult that co-owns
+	//! the context (so it survives `del conn`). Never call these on a StreamQueryResult:
+	//! a lazy result already has a live context and is converted/wrapped directly.
+	void PromoteMaterializedToArrow(idx_t batch_size);
+
+	template <typename T>
+	T RunWithArrowSchema(const std::function<T(const ArrowSchema &)> &fun, bool dedup_col_names);
+	duckdb::pyarrow::Table MaterializedResultToArrowTable(const ArrowSchema &arrow_schema, idx_t rows_per_batch);
+	ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch);
+
 private:
 	idx_t chunk_offset = 0;
 

src/duckdb_py/pyrelation.cpp

@@ -960,11 +960,11 @@ PandasDataFrame DuckDBPyRelation::FetchDFChunk(idx_t vectors_per_chunk, bool dat
 	return result->FetchDFChunk(vectors_per_chunk, date_as_object);
 }
 
-duckdb::pyarrow::Table DuckDBPyRelation::ToArrowTableInternal(idx_t batch_size, bool to_polars) {
+pyarrow::Table DuckDBPyRelation::ToArrowTableInternal(idx_t batch_size, bool to_polars) {
+	if (!result && !rel) {
+		return py::none();
+	}
 	if (!result) {
-		if (!rel) {
-			return py::none();
-		}
 		auto &config = ClientConfig::GetConfig(*rel->context->GetContext());
 		ScopedConfigSetting scoped_setting(
 		    config,
@@ -991,19 +991,9 @@ py::object DuckDBPyRelation::ToArrowCapsule(const py::object &requested_schema)
 		if (!rel) {
 			return py::none();
 		}
-		// The PyCapsule protocol doesn't allow custom parameters, so we use the same
-		// default batch size as fetch_arrow_table / fetch_record_batch.
-		idx_t batch_size = 1000000;
-		auto &config = ClientConfig::GetConfig(*rel->context->GetContext());
-		ScopedConfigSetting scoped_setting(
-		    config,
-		    [&batch_size](ClientConfig &config) {
-			    config.get_result_collector = [&batch_size](ClientContext &context, PreparedStatementData &data) {
-				    return PhysicalArrowCollector::Create(context, data, batch_size);
-			    };
-		    },
-		    [](ClientConfig &config) { config.get_result_collector = nullptr; });
-		ExecuteOrThrow();
+		// Fresh relation: stream lazily on the user's context (capsule survives `del conn`,
+		// but shares the single active-stream slot - consume before reusing the connection).
+		ExecuteOrThrow(true);
 	}
 	AssertResultOpen();
 	auto res = result->FetchArrowCapsule();
@@ -1049,6 +1039,7 @@ duckdb::pyarrow::RecordBatchReader DuckDBPyRelation::ToRecordBatch(idx_t batch_s
 		if (!rel) {
 			return py::none();
 		}
+		// Fresh relation: stream lazily on the user's own context (survives `del conn`).
 		ExecuteOrThrow(true);
 	}
 	AssertResultOpen();