Skip to content

Commit 9e2e5a2

Browse files
committed
Merge branch 'v1.5-variegata' into main
Periodic forward-merge of release-branch bugfixes into main. Notably brings in #495 "Unify arrow exports across all query result types" (the materialized slow-path lifetime / connection-GC fix and the test_arrow_refeed suite), replacing main's older SchemaCachingStreamWrapper/ArrowQueryResultStreamWrapper approach. Submodule: external/duckdb is kept at main's pin 0361de441a (v1.5's submodule bumps discarded; git fast-forwarded the gitlink to main's newer pin). Conflict resolution: - .github/workflows/packaging_wheels.yml: applied both intents — v1.5's windows-2025 -> windows-2022 (consistent with targeted_test.yml) and main's ARM64-comment removal. Adaptation for main's newer core: - pyresult.cpp: core's ColumnDataRef now takes vector<Identifier> (not vector<string>); promote the deduplicated scan names to Identifiers explicitly in MakeColumnDataScanStatement. Verified: clean build; tests/fast/arrow + tests/fast/udf = 2436 passed, 0 failed (incl. test_capsule_slow_path_survives_connection_gc and the new test_arrow_refeed suite).
2 parents 01cf1f1 + c400b90 commit 9e2e5a2

11 files changed

Lines changed: 577 additions & 365 deletions

File tree

.github/workflows/packaging_wheels.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ jobs:
3232
matrix:
3333
python: [ cp314 ]
3434
platform:
35-
- { os: windows-2025, arch: amd64, cibw_system: win }
35+
- { os: windows-2022, arch: amd64, cibw_system: win }
3636
- { os: windows-11-arm, arch: ARM64, cibw_system: win }
3737
- { os: ubuntu-24.04, arch: x86_64, cibw_system: manylinux }
3838
- { os: ubuntu-24.04-arm, arch: aarch64, cibw_system: manylinux }

.github/workflows/targeted_test.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ on:
88
required: true
99
type: choice
1010
options:
11-
- 'windows-2025'
11+
- 'windows-2022'
1212
- 'windows-11-arm'
1313
- 'ubuntu-24.04'
1414
- 'ubuntu-24.04-arm'

src/duckdb_py/arrow/arrow_array_stream.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,11 @@
1414

1515
namespace duckdb {
1616

17-
void TransformDuckToArrowChunk(ArrowSchema &arrow_schema, ArrowArray &data, py::list &batches) {
17+
void TransformDuckToArrowChunk(py::object pyarrow_schema, ArrowArray &data, py::list &batches) {
1818
py::gil_assert();
1919
auto pyarrow_lib_module = py::module::import("pyarrow").attr("lib");
2020
auto batch_import_func = pyarrow_lib_module.attr("RecordBatch").attr("_import_from_c");
21-
batches.append(batch_import_func(reinterpret_cast<uint64_t>(&data), reinterpret_cast<uint64_t>(&arrow_schema)));
21+
batches.append(batch_import_func(reinterpret_cast<uint64_t>(&data), pyarrow_schema));
2222
}
2323

2424
void VerifyArrowDatasetLoaded() {

src/duckdb_py/arrow/arrow_export_utils.cpp

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,18 +17,28 @@ namespace duckdb {
1717

1818
namespace pyarrow {
1919

20-
py::object ToArrowTable(const vector<LogicalType> &types, const vector<string> &names, const py::list &batches,
21-
ClientProperties &options) {
20+
py::object ToPyArrowSchema(const ArrowSchema &schema) {
2221
py::gil_scoped_acquire acquire;
2322

2423
auto pyarrow_lib_module = py::module::import("pyarrow").attr("lib");
25-
auto from_batches_func = pyarrow_lib_module.attr("Table").attr("from_batches");
2624
auto schema_import_func = pyarrow_lib_module.attr("Schema").attr("_import_from_c");
25+
return schema_import_func(reinterpret_cast<uint64_t>(&schema));
26+
}
27+
28+
py::object ToArrowTable(const py::list &batches, py::object pyarrow_schema) {
29+
py::gil_scoped_acquire acquire;
30+
31+
auto pyarrow_lib_module = py::module::import("pyarrow").attr("lib");
32+
auto from_batches_func = pyarrow_lib_module.attr("Table").attr("from_batches");
33+
34+
return py::cast<duckdb::pyarrow::Table>(from_batches_func(batches, pyarrow_schema));
35+
}
36+
37+
py::object ToArrowTable(const vector<LogicalType> &types, const vector<string> &names, const py::list &batches,
38+
ClientProperties &options) {
2739
ArrowSchema schema;
2840
ArrowConverter::ToArrowSchema(&schema, types, names, options);
29-
auto schema_obj = schema_import_func(reinterpret_cast<uint64_t>(&schema));
30-
31-
return py::cast<duckdb::pyarrow::Table>(from_batches_func(batches, schema_obj));
41+
return ToArrowTable(batches, ToPyArrowSchema(schema));
3242
}
3343

3444
} // namespace pyarrow

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ enum class PyArrowObjectType {
6262
PolarsLazyFrame
6363
};
6464

65-
void TransformDuckToArrowChunk(ArrowSchema &arrow_schema, ArrowArray &data, py::list &batches);
65+
void TransformDuckToArrowChunk(py::object pyarrow_schema, ArrowArray &data, py::list &batches);
6666

6767
PyArrowObjectType GetArrowType(const py::handle &obj);
6868

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,13 @@ namespace duckdb {
66

77
namespace pyarrow {
88

9+
py::object ToPyArrowSchema(const ArrowSchema &schema);
10+
911
py::object ToArrowTable(const vector<LogicalType> &types, const vector<string> &names, const py::list &batches,
1012
ClientProperties &options);
1113

14+
py::object ToArrowTable(const py::list &batches, py::object pyarrow_schema);
15+
1216
} // namespace pyarrow
1317

1418
} // namespace duckdb

src/duckdb_py/include/duckdb_python/pyresult.hpp

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,15 +36,13 @@ struct DuckDBPyResult {
3636

3737
PandasDataFrame FetchDF(bool date_as_object);
3838

39-
duckdb::pyarrow::Table FetchArrowTable(idx_t rows_per_batch, bool to_polars);
40-
4139
PandasDataFrame FetchDFChunk(const idx_t vectors_per_chunk = 1, bool date_as_object = false);
4240

4341
py::dict FetchPyTorch();
4442

4543
py::dict FetchTF();
4644

47-
ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch = 1000000);
45+
duckdb::pyarrow::Table FetchArrowTable(idx_t rows_per_batch, bool to_polars);
4846
duckdb::pyarrow::RecordBatchReader FetchRecordBatchReader(idx_t rows_per_batch = 1000000);
4947
py::object FetchArrowCapsule(idx_t rows_per_batch = 1000000);
5048

@@ -71,6 +69,19 @@ struct DuckDBPyResult {
7169
unique_ptr<DataChunk> FetchNextRaw(QueryResult &result);
7270
unique_ptr<NumpyResultConversion> InitializeNumpyConversion(bool pandas = false);
7371

72+
//! Re-feed an already-MATERIALIZED result (a ColumnDataCollection, e.g. from
73+
//! rel.execute()) back through the engine on the user's own context. The eager
74+
//! variant installs a PhysicalArrowCollector to produce an ArrowQueryResult
75+
//! (parallel); the stream variant produces a lazy StreamQueryResult that co-owns
76+
//! the context (so it survives `del conn`). Never call these on a StreamQueryResult:
77+
//! a lazy result already has a live context and is converted/wrapped directly.
78+
void PromoteMaterializedToArrow(idx_t batch_size);
79+
80+
template <typename T>
81+
T RunWithArrowSchema(const std::function<T(const ArrowSchema &)> &fun, bool dedup_col_names);
82+
duckdb::pyarrow::Table MaterializedResultToArrowTable(const ArrowSchema &arrow_schema, idx_t rows_per_batch);
83+
ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch);
84+
7485
private:
7586
idx_t chunk_offset = 0;
7687

src/duckdb_py/pyrelation.cpp

Lines changed: 8 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -960,11 +960,11 @@ PandasDataFrame DuckDBPyRelation::FetchDFChunk(idx_t vectors_per_chunk, bool dat
960960
return result->FetchDFChunk(vectors_per_chunk, date_as_object);
961961
}
962962

963-
duckdb::pyarrow::Table DuckDBPyRelation::ToArrowTableInternal(idx_t batch_size, bool to_polars) {
963+
pyarrow::Table DuckDBPyRelation::ToArrowTableInternal(idx_t batch_size, bool to_polars) {
964+
if (!result && !rel) {
965+
return py::none();
966+
}
964967
if (!result) {
965-
if (!rel) {
966-
return py::none();
967-
}
968968
auto &config = ClientConfig::GetConfig(*rel->context->GetContext());
969969
ScopedConfigSetting scoped_setting(
970970
config,
@@ -991,19 +991,9 @@ py::object DuckDBPyRelation::ToArrowCapsule(const py::object &requested_schema)
991991
if (!rel) {
992992
return py::none();
993993
}
994-
// The PyCapsule protocol doesn't allow custom parameters, so we use the same
995-
// default batch size as fetch_arrow_table / fetch_record_batch.
996-
idx_t batch_size = 1000000;
997-
auto &config = ClientConfig::GetConfig(*rel->context->GetContext());
998-
ScopedConfigSetting scoped_setting(
999-
config,
1000-
[&batch_size](ClientConfig &config) {
1001-
config.get_result_collector = [&batch_size](ClientContext &context, PreparedStatementData &data) {
1002-
return PhysicalArrowCollector::Create(context, data, batch_size);
1003-
};
1004-
},
1005-
[](ClientConfig &config) { config.get_result_collector = nullptr; });
1006-
ExecuteOrThrow();
994+
// Fresh relation: stream lazily on the user's context (capsule survives `del conn`,
995+
// but shares the single active-stream slot - consume before reusing the connection).
996+
ExecuteOrThrow(true);
1007997
}
1008998
AssertResultOpen();
1009999
auto res = result->FetchArrowCapsule();
@@ -1049,6 +1039,7 @@ duckdb::pyarrow::RecordBatchReader DuckDBPyRelation::ToRecordBatch(idx_t batch_s
10491039
if (!rel) {
10501040
return py::none();
10511041
}
1042+
// Fresh relation: stream lazily on the user's own context (survives `del conn`).
10521043
ExecuteOrThrow(true);
10531044
}
10541045
AssertResultOpen();

0 commit comments

Comments
 (0)