diff --git a/external/duckdb b/external/duckdb index 81bc275d..4d95776f 160000 --- a/external/duckdb +++ b/external/duckdb @@ -1 +1 @@ -Subproject commit 81bc275dd617edf26d402d712859c5ccd9884e04 +Subproject commit 4d95776fd7774793d7d4a8d47eca9f6d9089da09 diff --git a/pyproject.toml b/pyproject.toml index 128f0ecd..bf23ab4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -327,6 +327,10 @@ build = [ "cmake>=3.29.0", "ninja>=1.10", "nanobind>=3.0", + # C API headers for the editable build. Unpinned because the test group forces numpy<2 on + # Python 3.11 for tensorflow, and one universal resolution must satisfy both. Release wheels + # build isolated against the numpy>=2.0 in [build-system]. + "numpy", "scikit_build_core>=0.11.4", ] dev = [ # tooling like uv will install this automatically when syncing the environment diff --git a/src/arrow/arrow_array_stream.cpp b/src/arrow/arrow_array_stream.cpp index 99e02b46..5cbc4f68 100644 --- a/src/arrow/arrow_array_stream.cpp +++ b/src/arrow/arrow_array_stream.cpp @@ -33,7 +33,7 @@ nb::object PythonTableArrowArrayStreamFactory::ProduceScanner(nb::object &arrow_ const ClientProperties &client_properties) { D_ASSERT(!nb::isinstance(arrow_obj_handle)); ArrowSchemaWrapper schema; - PythonTableArrowArrayStreamFactory::GetSchemaInternal(arrow_obj_handle, schema); + PythonTableArrowArrayStreamFactory::GetSchemaInternal(arrow_obj_handle, schema.arrow_schema); ArrowTableSchema arrow_table; ArrowTableFunction::PopulateArrowTableSchema(*client_properties.client_context.get_mutable(), arrow_table, schema.arrow_schema); @@ -59,13 +59,12 @@ nb::object PythonTableArrowArrayStreamFactory::ProduceScanner(nb::object &arrow_ return arrow_scanner(arrow_obj_handle, **kwargs); } -unique_ptr PythonTableArrowArrayStreamFactory::Produce(uintptr_t factory_ptr, - ArrowStreamParameters ¶meters) { +unique_ptr +PythonTableArrowArrayStreamFactory::ProduceStream(ArrowStreamParameters ¶meters) { nb::gil_scoped_acquire acquire; - auto factory = static_cast(reinterpret_cast(factory_ptr)); // NOLINT - D_ASSERT(factory->arrow_object); - nb::handle arrow_obj_handle(factory->arrow_object); - auto arrow_object_type = factory->cached_arrow_type; + D_ASSERT(arrow_object.obj.ptr()); + nb::handle arrow_obj_handle(arrow_object.obj); + auto arrow_object_type = cached_arrow_type; if (arrow_object_type == PyArrowObjectType::PolarsLazyFrame) { nb::object lf = nb::borrow(arrow_obj_handle); @@ -81,9 +80,9 @@ unique_ptr PythonTableArrowArrayStreamFactory::Produce( // rather than silently returning unfiltered rows — the arrow scan does not // re-apply pushed filters. Mirrors the pyarrow ProduceScanner path. if (filters && filters->HasFilters()) { - auto filter_expr = PolarsFilterPushdown::TransformFilter( - *filters, parameters.projected_columns.projection_map, parameters.projected_columns.filter_to_col, - factory->client_properties); + auto filter_expr = + PolarsFilterPushdown::TransformFilter(*filters, parameters.projected_columns.projection_map, + parameters.projected_columns.filter_to_col, client_properties); if (!filter_expr.is(nb::none())) { lf = lf.attr("filter")(filter_expr); filters_pushed = true; @@ -93,13 +92,13 @@ unique_ptr PythonTableArrowArrayStreamFactory::Produce( // If no filters were pushed and we have a cached Arrow table, reuse it. This avoids re-reading from source and // re-converting on repeated unfiltered scans. nb::object arrow_table; - if (!filters_pushed && factory->cached_arrow_table.ptr() != nullptr) { - arrow_table = factory->cached_arrow_table; + if (!filters_pushed && cached_arrow_table.obj.ptr() != nullptr) { + arrow_table = cached_arrow_table.obj; } else { arrow_table = lf.attr("collect")().attr("to_arrow")(); // Cache only unfiltered results (filtered results are partial) if (!filters_pushed) { - factory->cached_arrow_table = arrow_table; + cached_arrow_table.obj = arrow_table; } } @@ -141,7 +140,7 @@ unique_ptr PythonTableArrowArrayStreamFactory::Produce( auto &import_cache = *DuckDBPyConnection::ImportCache(); nb::object arrow_batch_scanner = import_cache.pyarrow.dataset.Scanner().attr("from_batches"); nb::handle reader_handle = reader; - auto scanner = ProduceScanner(arrow_batch_scanner, reader_handle, parameters, factory->client_properties); + auto scanner = ProduceScanner(arrow_batch_scanner, reader_handle, parameters, client_properties); auto record_batches = scanner.attr("to_reader")(); auto res = make_uniq(); auto export_to_c = record_batches.attr("_export_to_c"); @@ -179,12 +178,12 @@ unique_ptr PythonTableArrowArrayStreamFactory::Produce( // If it's a scanner we have to turn it to a record batch reader, and then a scanner again since we can't stack // scanners on arrow Otherwise pushed-down projections and filters will disappear like tears in the rain auto record_batches = arrow_obj_handle.attr("to_reader")(); - scanner = ProduceScanner(arrow_batch_scanner, record_batches, parameters, factory->client_properties); + scanner = ProduceScanner(arrow_batch_scanner, record_batches, parameters, client_properties); break; } case PyArrowObjectType::Dataset: { nb::object arrow_scanner = arrow_obj_handle.attr("__class__").attr("scanner"); - scanner = ProduceScanner(arrow_scanner, arrow_obj_handle, parameters, factory->client_properties); + scanner = ProduceScanner(arrow_scanner, arrow_obj_handle, parameters, client_properties); break; } default: { @@ -201,7 +200,7 @@ unique_ptr PythonTableArrowArrayStreamFactory::Produce( return res; } -void PythonTableArrowArrayStreamFactory::GetSchemaInternal(nb::handle arrow_obj_handle, ArrowSchemaWrapper &schema) { +void PythonTableArrowArrayStreamFactory::GetSchemaInternal(nb::handle arrow_obj_handle, ArrowSchema &schema) { // PyCapsule (from bare capsule Produce path) if (nb::isinstance(arrow_obj_handle)) { auto capsule = nb::borrow(arrow_obj_handle); @@ -209,7 +208,7 @@ void PythonTableArrowArrayStreamFactory::GetSchemaInternal(nb::handle arrow_obj_ if (!stream->release) { throw InvalidInputException("This ArrowArrayStream has already been consumed and cannot be scanned again."); } - if (stream->get_schema(stream, &schema.arrow_schema)) { + if (stream->get_schema(stream, &schema)) { throw InvalidInputException("Failed to get Arrow schema from stream: %s", stream->get_last_error ? stream->get_last_error(stream) : "unknown error"); } @@ -221,28 +220,25 @@ void PythonTableArrowArrayStreamFactory::GetSchemaInternal(nb::handle arrow_obj_ auto &import_cache = *DuckDBPyConnection::ImportCache(); if (duckdb::PyUtil::IsInstance(arrow_obj_handle, import_cache.pyarrow.dataset.Scanner())) { auto obj_schema = arrow_obj_handle.attr("projected_schema"); - obj_schema.attr("_export_to_c")(reinterpret_cast(&schema.arrow_schema)); + obj_schema.attr("_export_to_c")(reinterpret_cast(&schema)); } else { auto obj_schema = arrow_obj_handle.attr("schema"); - obj_schema.attr("_export_to_c")(reinterpret_cast(&schema.arrow_schema)); + obj_schema.attr("_export_to_c")(reinterpret_cast(&schema)); } } -void PythonTableArrowArrayStreamFactory::GetSchema(uintptr_t factory_ptr, ArrowSchemaWrapper &schema) { - auto factory = static_cast(reinterpret_cast(factory_ptr)); // NOLINT - - // Fast path: return cached schema without GIL or Python calls - if (factory->schema_cached) { - schema.arrow_schema = factory->cached_schema; // struct copy - schema.arrow_schema.release = nullptr; // non-owning copy +void PythonTableArrowArrayStreamFactory::GetSchema(ArrowSchema &schema) { + if (schema_cached.load(std::memory_order_acquire)) { + schema = cached_schema; + schema.release = nullptr; return; } nb::gil_scoped_acquire acquire; - D_ASSERT(factory->arrow_object); - nb::handle arrow_obj_handle(factory->arrow_object); + D_ASSERT(arrow_object.obj.ptr()); + nb::handle arrow_obj_handle(arrow_object.obj); - auto type = factory->cached_arrow_type; + auto type = cached_arrow_type; if (type == PyArrowObjectType::PolarsLazyFrame) { // head(0).collect().to_arrow() gives the Arrow-exported schema (e.g. large_string) without materializing data. // collect_schema() would give Polars-native types (e.g. string_view) that don't match the actual export. @@ -250,11 +246,11 @@ void PythonTableArrowArrayStreamFactory::GetSchema(uintptr_t factory_ptr, ArrowS const auto schema_capsule = empty_arrow.attr("schema").attr("__arrow_c_schema__")(); const auto capsule = nb::borrow(schema_capsule); const auto arrow_schema = reinterpret_cast(capsule.data()); - factory->cached_schema = *arrow_schema; + cached_schema = *arrow_schema; arrow_schema->release = nullptr; - factory->schema_cached = true; - schema.arrow_schema = factory->cached_schema; - schema.arrow_schema.release = nullptr; + schema_cached.store(true, std::memory_order_release); + schema = cached_schema; + schema.release = nullptr; return; } if (type == PyArrowObjectType::PyCapsuleInterface || type == PyArrowObjectType::Table) { @@ -263,18 +259,18 @@ void PythonTableArrowArrayStreamFactory::GetSchema(uintptr_t factory_ptr, ArrowS auto schema_capsule = arrow_obj_handle.attr("__arrow_c_schema__")(); auto capsule = nb::borrow(schema_capsule); auto arrow_schema = reinterpret_cast(capsule.data()); - factory->cached_schema = *arrow_schema; // factory takes ownership + cached_schema = *arrow_schema; arrow_schema->release = nullptr; - factory->schema_cached = true; - schema.arrow_schema = factory->cached_schema; // non-owning copy - schema.arrow_schema.release = nullptr; + schema_cached.store(true, std::memory_order_release); + schema = cached_schema; + schema.release = nullptr; return; } // Otherwise try to use .schema with _export_to_c if (nb::hasattr(arrow_obj_handle, "schema")) { auto obj_schema = arrow_obj_handle.attr("schema"); if (nb::hasattr(obj_schema, "_export_to_c")) { - obj_schema.attr("_export_to_c")(reinterpret_cast(&schema.arrow_schema)); + obj_schema.attr("_export_to_c")(reinterpret_cast(&schema)); return; } } @@ -282,7 +278,7 @@ void PythonTableArrowArrayStreamFactory::GetSchema(uintptr_t factory_ptr, ArrowS auto stream_capsule = arrow_obj_handle.attr("__arrow_c_stream__")(); auto capsule = nb::borrow(stream_capsule); auto stream = reinterpret_cast(capsule.data()); - if (stream->get_schema(stream, &schema.arrow_schema)) { + if (stream->get_schema(stream, &schema)) { throw InvalidInputException("Failed to get Arrow schema from stream: %s", stream->get_last_error ? stream->get_last_error(stream) : "unknown error"); } @@ -292,9 +288,9 @@ void PythonTableArrowArrayStreamFactory::GetSchema(uintptr_t factory_ptr, ArrowS // Cache for Table and Dataset (immutable schema) if (type == PyArrowObjectType::Table || type == PyArrowObjectType::Dataset) { - factory->cached_schema = schema.arrow_schema; // factory takes ownership - schema.arrow_schema.release = nullptr; // caller gets non-owning copy - factory->schema_cached = true; + cached_schema = schema; + schema.release = nullptr; + schema_cached.store(true, std::memory_order_release); } } diff --git a/src/include/duckdb_python/arrow/arrow_array_stream.hpp b/src/include/duckdb_python/arrow/arrow_array_stream.hpp index 39566beb..17aaf7d6 100644 --- a/src/include/duckdb_python/arrow/arrow_array_stream.hpp +++ b/src/include/duckdb_python/arrow/arrow_array_stream.hpp @@ -16,6 +16,7 @@ #include "duckdb/main/client_config.hpp" #include "duckdb/main/config.hpp" #include "duckdb_python/nb/casters.hpp" +#include "duckdb_python/registered_py_object.hpp" #include "duckdb/common/string.hpp" #include "duckdb/common/vector.hpp" @@ -66,44 +67,43 @@ void TransformDuckToArrowChunk(nb::object pyarrow_schema, ArrowArray &data, nb:: PyArrowObjectType GetArrowType(const nb::handle &obj); -class PythonTableArrowArrayStreamFactory { +class PythonTableArrowArrayStreamFactory : public ArrowScanFactory { public: - explicit PythonTableArrowArrayStreamFactory(PyObject *arrow_table, const ClientProperties &client_properties_p, - PyArrowObjectType arrow_type_p) - : arrow_object(arrow_table), client_properties(client_properties_p), cached_arrow_type(arrow_type_p) { + //! Must be constructed while holding the GIL. + PythonTableArrowArrayStreamFactory(nb::object arrow_object_p, const ClientProperties &client_properties_p, + PyArrowObjectType arrow_type_p) + : arrow_object(std::move(arrow_object_p)), client_properties(client_properties_p), + cached_arrow_type(arrow_type_p) { cached_schema.release = nullptr; } - ~PythonTableArrowArrayStreamFactory() { - if (cached_arrow_table.ptr() != nullptr) { - nb::gil_scoped_acquire acquire; - cached_arrow_table = nb::object(); - } + ~PythonTableArrowArrayStreamFactory() override { + // The release callback of a schema taken from a Python producer may itself be Python. if (cached_schema.release) { - cached_schema.release(&cached_schema); + if (nb::detail::cleanup_guard guard {}) { + cached_schema.release(&cached_schema); + } } } - //! Produces an Arrow Scanner, should be only called once when initializing Scan States - static unique_ptr Produce(uintptr_t factory, ArrowStreamParameters ¶meters); + void GetSchema(ArrowSchema &schema) override; + unique_ptr ProduceStream(ArrowStreamParameters ¶meters) override; - //! Get the schema of the arrow object - static void GetSchemaInternal(nb::handle arrow_object, ArrowSchemaWrapper &schema); - static void GetSchema(uintptr_t factory_ptr, ArrowSchemaWrapper &schema); + static void GetSchemaInternal(nb::handle arrow_object, ArrowSchema &schema); //! Arrow Object (i.e., Scanner, Record Batch Reader, Table, Dataset) - PyObject *arrow_object; + PyObjectHolder arrow_object; const ClientProperties client_properties; const PyArrowObjectType cached_arrow_type; //! Cached Arrow table from an unfiltered .collect().to_arrow() on a LazyFrame. //! Avoids re-reading from source and re-converting on repeated scans without filters. - nb::object cached_arrow_table; + PyObjectHolder cached_arrow_table; private: ArrowSchema cached_schema; - bool schema_cached = false; + atomic schema_cached {false}; static nb::object ProduceScanner(nb::object &arrow_scanner, nb::handle &arrow_obj_handle, ArrowStreamParameters ¶meters, const ClientProperties &client_properties); diff --git a/src/include/duckdb_python/filesystem_object.hpp b/src/include/duckdb_python/filesystem_object.hpp index e35da40f..ddec02ea 100644 --- a/src/include/duckdb_python/filesystem_object.hpp +++ b/src/include/duckdb_python/filesystem_object.hpp @@ -18,8 +18,10 @@ class FileSystemObject : public RegisteredObject { : RegisteredObject(std::move(fs)), filenames(std::move(filenames_p)) { } ~FileSystemObject() override { - nb::gil_scoped_acquire acquire; - // Assert that the 'obj' is a filesystem + nb::detail::cleanup_guard guard {}; + if (!guard) { + return; + } D_ASSERT(duckdb::PyUtil::IsInstance( obj, DuckDBPyConnection::ImportCache()->duckdb.filesystem.ModifiedMemoryFileSystem())); // Destructors are implicitly noexcept: a Python exception escaping here (fsspec `_rm` raises diff --git a/src/include/duckdb_python/map.hpp b/src/include/duckdb_python/map.hpp index c1b6325b..68a6d8ab 100644 --- a/src/include/duckdb_python/map.hpp +++ b/src/include/duckdb_python/map.hpp @@ -10,11 +10,21 @@ #include "duckdb.hpp" #include "duckdb_python/nb/casters.hpp" +#include "duckdb_python/registered_py_object.hpp" #include "duckdb/parser/parsed_data/create_table_function_info.hpp" #include "duckdb/execution/execution_context.hpp" namespace duckdb { +//! Carried through the bind input so that SQL text can never forge a reference to the callable +struct MapFunctionInfo : public TableFunctionInfo { + MapFunctionInfo(nb::object function_p, nb::object schema_p) + : function(std::move(function_p)), schema(std::move(schema_p)) { + } + PyObjectHolder function; + PyObjectHolder schema; +}; + struct MapFunction : public TableFunction { public: diff --git a/src/include/duckdb_python/pandas/pandas_scan.hpp b/src/include/duckdb_python/pandas/pandas_scan.hpp index b0356644..1b74a06d 100644 --- a/src/include/duckdb_python/pandas/pandas_scan.hpp +++ b/src/include/duckdb_python/pandas/pandas_scan.hpp @@ -13,9 +13,17 @@ #include "duckdb_python/pandas/pandas_bind.hpp" #include "duckdb_python/nb/casters.hpp" +#include "duckdb_python/registered_py_object.hpp" namespace duckdb { +//! Carried through the bind input so that SQL text can never name the dataframe +struct PandasScanInfo : public TableFunctionInfo { + explicit PandasScanInfo(nb::object df_p) : df(std::move(df_p)) { + } + PyObjectHolder df; +}; + struct PandasScanFunction : public TableFunction { public: static constexpr idx_t PANDAS_PARTITION_COUNT = 50 * STANDARD_VECTOR_SIZE; diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 638b0a4b..67e0fd8c 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -31,14 +31,6 @@ enum class PythonEnvironmentType { NORMAL, INTERACTIVE, JUPYTER }; struct DuckDBPyRelation; -class RegisteredArrow : public RegisteredObject { - -public: - RegisteredArrow(unique_ptr arrow_factory_p, nb::object obj_p) - : RegisteredObject(std::move(obj_p)), arrow_factory(std::move(arrow_factory_p)) {}; - unique_ptr arrow_factory; -}; - struct DefaultConnectionHolder { public: DefaultConnectionHolder() { @@ -185,7 +177,7 @@ struct DuckDBPyConnection : public std::enable_shared_from_this> statements); unique_ptr PrepareQuery(unique_ptr statement); unique_ptr ExecuteInternal(PreparedStatement &prep, nb::object params = nb::list()); - unique_ptr PrepareAndExecuteInternal(unique_ptr statement, - nb::object params = nb::list()); + //! Binds the parameters and submits the statement. The handle is returned undriven. + unique_ptr PrepareAndSubmitInternal(unique_ptr statement, + nb::object params = nb::list()); std::shared_ptr Execute(const nb::object &query, nb::object params = nb::list()); std::shared_ptr ExecuteFromString(const string &query); @@ -367,7 +360,9 @@ struct DuckDBPyConnection : public std::enable_shared_from_this CompletePendingQuery(PendingQueryResult &pending_query); + //! Runs a submitted query to a retained, ended result on the calling thread, checking for Python + //! signals between tasks. Throws the query's error, and leaves the handle holding its collection. + static void CompleteQuery(QueryResult &result); private: std::unique_ptr CreateRelation(shared_ptr rel); diff --git a/src/include/duckdb_python/pyrelation.hpp b/src/include/duckdb_python/pyrelation.hpp index f71a6327..2b27fceb 100644 --- a/src/include/duckdb_python/pyrelation.hpp +++ b/src/include/duckdb_python/pyrelation.hpp @@ -186,7 +186,7 @@ struct DuckDBPyRelation { nb::dict FetchTF(); - nb::dict FetchNumpyInternal(bool stream = false, idx_t vectors_per_chunk = 1); + nb::dict FetchNumpyInternal(bool chunked = false, idx_t vectors_per_chunk = 1); PandasDataFrame FetchDFChunk(const idx_t vectors_per_chunk = 1, bool date_as_object = false); @@ -287,7 +287,8 @@ struct DuckDBPyRelation { void AssertResultOpen() const; void AssertRelation() const; void ExecuteOrThrow(bool stream_result = false); - unique_ptr ExecuteInternal(bool stream_result = false); + //! Runs the relation to a completed, retained result + unique_ptr ExecuteInternal(); private: //! Prevents GC of the parent DuckDBPyConnection. diff --git a/src/include/duckdb_python/pyresult.hpp b/src/include/duckdb_python/pyresult.hpp index 475b49c8..d9b56d19 100644 --- a/src/include/duckdb_python/pyresult.hpp +++ b/src/include/duckdb_python/pyresult.hpp @@ -10,7 +10,7 @@ #include "duckdb_python/numpy/numpy_result_conversion.hpp" #include "duckdb.hpp" -#include "duckdb/main/chunk_scan_state.hpp" +#include "duckdb/main/query_result_stream.hpp" #include "duckdb_python/nb/casters.hpp" #include "duckdb_python/python_objects.hpp" #include "duckdb_python/dataframe.hpp" @@ -19,7 +19,11 @@ namespace duckdb { struct DuckDBPyResult { public: - explicit DuckDBPyResult(unique_ptr result); + //! A result that has run to completion: it holds its rows, or is an Arrow result + explicit DuckDBPyResult(unique_ptr completed); + //! A freshly submitted handle. Opened as a stream when the caller asks for one and the statement + //! can be drained, otherwise run to completion. Call with the GIL released. + DuckDBPyResult(unique_ptr submitted, bool stream_result); ~DuckDBPyResult(); public: @@ -31,7 +35,7 @@ struct DuckDBPyResult { nb::dict FetchNumpy(); - nb::dict FetchNumpyInternal(bool stream = false, idx_t vectors_per_chunk = 1, + nb::dict FetchNumpyInternal(bool chunked = false, idx_t vectors_per_chunk = 1, std::unique_ptr conversion = nullptr); PandasDataFrame FetchDF(bool date_as_object); @@ -50,14 +54,12 @@ struct DuckDBPyResult { void Close(); - bool IsClosed() const; - unique_ptr FetchChunk(); vector GetNames(); - const vector &GetTypes(); + const vector &GetTypes() const; - ClientProperties GetClientProperties(); + const ClientProperties &GetClientProperties() const; private: void FillNumpy(nb::dict &res, idx_t col_idx, NumpyResultConversion &conversion, const char *name); @@ -67,27 +69,46 @@ struct DuckDBPyResult { void ConvertDateTimeTypes(PandasDataFrame &df, bool date_as_object) const; //! The names the Python layer reports, see the definition for why this is not always the result's own. const vector &ResultNames() const; - unique_ptr FetchNext(QueryResult &result); - unique_ptr FetchNextRaw(QueryResult &result); + bool Empty() const { + return !result && !stream && !submitted; + } + //! The Arrow readers take the handle itself, so the stream is not opened before the first + //! consumer call decides. + void EnsureStream(); + //! Flat vectors, for the row fetch + unique_ptr FetchNext(); + unique_ptr FetchNextRaw(); + unique_ptr FetchStreamChunk(); + //! Rows a row fetch popped but has not returned yet. Once the stream's query has ended + //! underneath them they are dropped, so that the engine reports why on the next fetch. + unique_ptr TakeBufferedRows(); + //! Takes the context lock: call with the GIL released, the engine acquires the GIL under it. + bool StreamEnded() const { + return stream && !stream->IsOpen(); + } + void Retain(); + void CloseStream(); std::unique_ptr 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. + //! Re-feed a retained result's collection through a PhysicalArrowCollector on the user's own + //! context, which converts in parallel and yields an ArrowQueryResult in its place. void PromoteMaterializedToArrow(idx_t batch_size); template T RunWithArrowSchema(const std::function &fun, bool dedup_col_names); duckdb::pyarrow::Table MaterializedResultToArrowTable(const ArrowSchema &arrow_schema, idx_t rows_per_batch); + //! The stream's private data owns the handle, and its callbacks run without the GIL. ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch); + bool IsArrow() const; private: idx_t chunk_offset = 0; - + //! The completed result, when the rows were retained or converted to Arrow unique_ptr result; + //! The open stream the rows are drained through + unique_ptr stream; + //! A handle that can be drained but has not been consumed yet + unique_ptr submitted; //! Set only when the result was re-bound (promotion to Arrow de-duplicates column names //! and core exposes no setter), so the original names survive. Empty means "use result's". vector names_override; @@ -96,7 +117,6 @@ struct DuckDBPyResult { unordered_map categories; // Holds the categorical type of Categorical/ENUM types unordered_map categories_type; - bool result_closed = false; }; } // namespace duckdb diff --git a/src/include/duckdb_python/python_dependency.hpp b/src/include/duckdb_python/python_dependency.hpp index 1531d25f..15add5f0 100644 --- a/src/include/duckdb_python/python_dependency.hpp +++ b/src/include/duckdb_python/python_dependency.hpp @@ -12,7 +12,6 @@ namespace duckdb { class PythonDependencyItem : public DependencyItem { public: explicit PythonDependencyItem(unique_ptr &&object); - ~PythonDependencyItem() override; public: static shared_ptr Create(nb::object object); diff --git a/src/include/duckdb_python/python_object_container.hpp b/src/include/duckdb_python/python_object_container.hpp index 60e3d716..4c585b75 100644 --- a/src/include/duckdb_python/python_object_container.hpp +++ b/src/include/duckdb_python/python_object_container.hpp @@ -22,8 +22,13 @@ class PythonObjectContainer { } ~PythonObjectContainer() { - nb::gil_scoped_acquire acquire; - py_obj.clear(); + if (nb::detail::cleanup_guard guard {}) { + py_obj.clear(); + } else { + for (auto &obj : py_obj) { + obj.release(); + } + } } void Push(nb::object &&obj) { diff --git a/src/include/duckdb_python/registered_py_object.hpp b/src/include/duckdb_python/registered_py_object.hpp index 809abfda..c6fe60c4 100644 --- a/src/include/duckdb_python/registered_py_object.hpp +++ b/src/include/duckdb_python/registered_py_object.hpp @@ -11,16 +11,35 @@ namespace duckdb { -class RegisteredObject { +//! Any engine thread may drop the reference, possibly while the interpreter is tearing down, in +//! which case it is leaked rather than touched. +class PyObjectHolder { public: - explicit RegisteredObject(nb::object obj_p) : obj(std::move(obj_p)) { + PyObjectHolder() = default; + //! Call with the GIL held + explicit PyObjectHolder(nb::object obj_p) : obj(std::move(obj_p)) { } - virtual ~RegisteredObject() { - nb::gil_scoped_acquire acquire; - obj = nb::none(); + PyObjectHolder(const PyObjectHolder &) = delete; + PyObjectHolder &operator=(const PyObjectHolder &) = delete; + PyObjectHolder(PyObjectHolder &&) = default; + //! Assignment would drop the previous reference outside the destructor's guard + PyObjectHolder &operator=(PyObjectHolder &&) = delete; + ~PyObjectHolder() { + if (nb::detail::cleanup_guard guard {}) { + obj = nb::object(); + } else { + obj.release(); + } } nb::object obj; }; +class RegisteredObject : public PyObjectHolder { +public: + explicit RegisteredObject(nb::object obj_p) : PyObjectHolder(std::move(obj_p)) { + } + virtual ~RegisteredObject() = default; +}; + } // namespace duckdb diff --git a/src/map.cpp b/src/map.cpp index 0347b19d..7b4ca6d5 100644 --- a/src/map.cpp +++ b/src/map.cpp @@ -1,4 +1,5 @@ #include "duckdb_python/map.hpp" +#include "duckdb/parser/tableref/table_function_ref.hpp" #include "duckdb_python/numpy/numpy_scan.hpp" #include "duckdb_python/pandas/pandas_bind.hpp" #include "duckdb_python/numpy/numpy_result_conversion.hpp" @@ -12,16 +13,14 @@ namespace duckdb { -MapFunction::MapFunction() - : TableFunction("python_map_function", {LogicalType::TABLE, LogicalType::POINTER, LogicalType::POINTER}, nullptr, - MapFunctionBind) { +MapFunction::MapFunction() : TableFunction("python_map_function", {LogicalType::TABLE}, nullptr, MapFunctionBind) { in_out_function = MapFunctionExec; } struct MapFunctionData : public TableFunctionData { - MapFunctionData() : function(nullptr) { + explicit MapFunctionData(shared_ptr info_p) : info(std::move(info_p)) { } - PyObject *function; + shared_ptr info; vector in_types, out_types; vector in_names, out_names; }; @@ -146,12 +145,17 @@ unique_ptr BindExplicitSchema(unique_ptr function // they better not change in the actual execution ^^ unique_ptr MapFunction::MapFunctionBind(ClientContext &context, TableFunctionBindInput &input, vector &return_types, vector &names) { + if (!input.ref.bind_info) { + throw BinderException("python_map_function requires a callable bind input"); + } + DynamicCastCheck(input.ref.bind_info.get()); + auto info = shared_ptr_cast(input.ref.bind_info); nb::gil_scoped_acquire acquire; - auto data_uptr = make_uniq(); + auto data_uptr = make_uniq(std::move(info)); auto &data = *data_uptr; - data.function = reinterpret_cast(input.inputs[1].GetPointer()); - auto explicit_schema = reinterpret_cast(input.inputs[2].GetPointer()); + auto function = data.info->function.obj.ptr(); + auto explicit_schema = data.info->schema.obj.ptr(); data.in_names = input.input_table_names; data.in_types = input.input_table_types; @@ -160,7 +164,7 @@ unique_ptr MapFunction::MapFunctionBind(ClientContext &context, Ta return BindExplicitSchema(std::move(data_uptr), explicit_schema, return_types, names); } NumpyResultConversion conversion(data.in_types, 0, context.GetClientProperties()); - auto df = FunctionCall(conversion, data.in_names, data.function); + auto df = FunctionCall(conversion, data.in_names, function); vector pandas_bind_data; // unused Pandas::Bind(context, df, pandas_bind_data, return_types, names); @@ -190,7 +194,7 @@ OperatorResultType MapFunction::MapFunctionExec(ExecutionContext &context, Table NumpyResultConversion conversion(data.in_types, input.size(), context.client.GetClientProperties()); conversion.Append(input); - auto df = FunctionCall(conversion, data.in_names, data.function); + auto df = FunctionCall(conversion, data.in_names, data.info->function.obj.ptr()); vector pandas_bind_data; vector pandas_return_types; diff --git a/src/pandas/scan.cpp b/src/pandas/scan.cpp index bc544a0c..e995896d 100644 --- a/src/pandas/scan.cpp +++ b/src/pandas/scan.cpp @@ -14,22 +14,27 @@ namespace duckdb { struct PandasScanFunctionData : public TableFunctionData { - PandasScanFunctionData(nb::handle df, idx_t row_count, vector pandas_bind_data, - vector sql_types, shared_ptr dependency) - : df(df), row_count(row_count), lines_read(0), pandas_bind_data(std::move(pandas_bind_data)), - sql_types(std::move(sql_types)), copied_df(std::move(dependency)) { + PandasScanFunctionData(shared_ptr info, idx_t row_count, + vector pandas_bind_data, vector sql_types) + : info(std::move(info)), row_count(row_count), lines_read(0), pandas_bind_data(std::move(pandas_bind_data)), + sql_types(std::move(sql_types)) { } - nb::handle df; + shared_ptr info; idx_t row_count; atomic lines_read; vector pandas_bind_data; vector sql_types; - shared_ptr copied_df; ~PandasScanFunctionData() override { try { - nb::gil_scoped_acquire acquire; - pandas_bind_data.clear(); + if (nb::detail::cleanup_guard guard {}) { + pandas_bind_data.clear(); + } else { + // Leaked on purpose: the Python references can no longer be dropped, and a static owner + // keeps a leak checker from reporting them + static auto *parked = new vector>(); + parked->push_back(std::move(pandas_bind_data)); + } } catch (...) { // NOLINT } } @@ -60,8 +65,7 @@ struct PandasScanGlobalState : public GlobalTableFunctionState { }; PandasScanFunction::PandasScanFunction() - : TableFunction("pandas_scan", {LogicalType::POINTER}, PandasScanFunc, PandasScanBind, PandasScanInitGlobal, - PandasScanInitLocal) { + : TableFunction("pandas_scan", {}, PandasScanFunc, PandasScanBind, PandasScanInitGlobal, PandasScanInitLocal) { get_partition_data = PandasScanGetPartitionData; cardinality = PandasScanCardinality; table_scan_progress = PandasProgress; @@ -82,8 +86,13 @@ OperatorPartitionData PandasScanFunction::PandasScanGetPartitionData(ClientConte unique_ptr PandasScanFunction::PandasScanBind(ClientContext &context, TableFunctionBindInput &input, vector &return_types, vector &names) { + if (!input.ref.bind_info) { + throw BinderException("pandas_scan requires a dataframe bind input"); + } + DynamicCastCheck(input.ref.bind_info.get()); + auto info = shared_ptr_cast(input.ref.bind_info); nb::gil_scoped_acquire acquire; - nb::handle df(reinterpret_cast(input.inputs[0].GetPointer())); + nb::handle df(info->df.obj); vector pandas_bind_data; @@ -94,22 +103,9 @@ unique_ptr PandasScanFunction::PandasScanBind(ClientContext &conte Pandas::Bind(context, df, pandas_bind_data, return_types, names); } auto df_columns = nb::list(df.attr("keys")()); - - auto &ref = input.ref; - - shared_ptr dependency_item; - if (ref.external_dependency) { - // This was created during the replacement scan if this was a pandas DataFrame (see python_replacement_scan.cpp) - dependency_item = ref.external_dependency->GetDependency("copy"); - if (!dependency_item) { - // This was created during the replacement if this was a numpy scan - dependency_item = ref.external_dependency->GetDependency("data"); - } - } - auto get_fun = df.attr("__getitem__"); idx_t row_count = nb::len(get_fun(df_columns[0])); - return make_uniq(df, row_count, std::move(pandas_bind_data), return_types, dependency_item); + return make_uniq(std::move(info), row_count, std::move(pandas_bind_data), return_types); } unique_ptr PandasScanFunction::PandasScanInitGlobal(ClientContext &context, diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index 5323996b..034f4a6b 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -36,11 +36,9 @@ #include "duckdb/function/function.hpp" #include "duckdb_python/nb/conversions/exception_handling_enum.hpp" #include "duckdb/parser/parsed_data/drop_info.hpp" -#include "duckdb/main/pending_query_result.hpp" +#include "duckdb/main/query_result.hpp" #include "duckdb_python/python_replacement_scan.hpp" #include "duckdb/common/shared_ptr.hpp" -#include "duckdb/main/materialized_query_result.hpp" -#include "duckdb/main/stream_query_result.hpp" #include "duckdb/main/relation/materialized_relation.hpp" #include "duckdb/parser/statement/load_statement.hpp" #include "duckdb_python/expression/pyexpression.hpp" @@ -554,26 +552,33 @@ std::shared_ptr DuckDBPyConnection::ExecuteMany(const nb::ob return shared_from_this(); } -unique_ptr DuckDBPyConnection::CompletePendingQuery(PendingQueryResult &pending_query) { - PendingExecutionResult execution_result; - if (pending_query.HasError()) { - pending_query.ThrowError(); +void DuckDBPyConnection::CompleteQuery(QueryResult &result) { + if (result.HasError()) { + result.ThrowError(); } - while (!PendingQueryResult::IsResultReady(execution_result = pending_query.ExecuteTask())) { + result.Materialize(); + // A result built by a delegating collector arrives finished and without a context; Poll reports + // that without touching the context, where ExecuteTask would throw. + auto state = result.Poll(); + while (!IsTerminal(state)) { { nb::gil_scoped_acquire gil; if (PyErr_CheckSignals() != 0) { throw std::runtime_error("Query interrupted"); } } - if (execution_result == PendingExecutionResult::BLOCKED) { - pending_query.WaitForTask(); + if (state == QueryResultState::BLOCKED || state == QueryResultState::READY || + state == QueryResultState::NO_TASKS_AVAILABLE) { + result.WaitForTask(); } + state = result.ExecuteTask(); } - if (execution_result == PendingExecutionResult::EXECUTION_ERROR) { - pending_query.ThrowError(); + // FINISHED only means the executor is done: the collection is taken and the query ended, with + // its transaction step, by Complete. Left open, the next statement would roll it back. + result.Complete(); + if (result.HasError()) { + result.ThrowError(); } - return pending_query.Execute(); } nb::list TransformNamedParameters(const case_insensitive_map_t &named_param_map, const nb::dict ¶ms) { @@ -670,21 +675,14 @@ unique_ptr DuckDBPyConnection::ExecuteInternal(PreparedStatement &p nb::gil_scoped_release release; unique_lock lock(py_connection_lock); - auto pending_query = prep.PendingQuery(named_values); - if (pending_query->HasError()) { - pending_query->ThrowError(); - } - res = CompletePendingQuery(*pending_query); - - if (res->HasError()) { - res->ThrowError(); - } + res = prep.Submit(named_values); + CompleteQuery(*res); } return res; } -unique_ptr DuckDBPyConnection::PrepareAndExecuteInternal(unique_ptr statement, - nb::object params) { +unique_ptr DuckDBPyConnection::PrepareAndSubmitInternal(unique_ptr statement, + nb::object params) { if (params.is_none()) { params = nb::list(); } @@ -698,14 +696,7 @@ unique_ptr DuckDBPyConnection::PrepareAndExecuteInternal(unique_ptr nb::gil_scoped_release release; unique_lock lock(py_connection_lock); - auto pending_query = con.GetConnection().PendingQuery(std::move(statement), named_values, true); - - if (pending_query->HasError()) { - pending_query->ThrowError(); - } - - res = CompletePendingQuery(*pending_query); - + res = con.GetConnection().Submit(std::move(statement), named_values); if (res->HasError()) { res->ThrowError(); } @@ -750,13 +741,19 @@ std::shared_ptr DuckDBPyConnection::Execute(const nb::object // FIXME: SQLites implementation says to not accept an 'execute' call with multiple statements ExecuteImmediately(std::move(statements)); - auto res = PrepareAndExecuteInternal(std::move(last_statement), std::move(params)); + auto res = PrepareAndSubmitInternal(std::move(last_statement), std::move(params)); - // Set the internal 'result' object if (res) { - // Don't use CreateRelation here — the result is stored inside the connection, - // so setting connection_owner would create a ref cycle (connection → result → connection). - con.SetResult(std::make_unique(std::make_shared(std::move(res)))); + std::shared_ptr py_result; + { + D_ASSERT(duckdb::PyUtil::GilCheck()); + nb::gil_scoped_release release; + unique_lock lock(py_connection_lock); + py_result = std::make_shared(std::move(res), true); + } + // Don't use CreateRelation here: the result is stored inside the connection, + // so setting connection_owner would create a ref cycle (connection, result, connection). + con.SetResult(std::make_unique(std::move(py_result))); } return shared_from_this(); } @@ -1624,15 +1621,8 @@ void DuckDBPyConnection::ExecuteImmediately(vector> sta "Prepared parameters are only supported for the last statement, please split your query up into " "separate 'execute' calls if you want to use prepared parameters"); } - auto pending_query = connection.PendingQuery(std::move(stmt), false); - if (pending_query->HasError()) { - pending_query->ThrowError(); - } - auto res = CompletePendingQuery(*pending_query); - - if (res->HasError()) { - res->ThrowError(); - } + auto res = connection.Submit(std::move(stmt)); + CompleteQuery(*res); } } @@ -1677,23 +1667,23 @@ std::unique_ptr DuckDBPyConnection::RunQuery(const nb::object if (!relation) { // Could not create a relation, resort to direct execution - unique_ptr res; - - res = PrepareAndExecuteInternal(std::move(last_statement), std::move(params)); - + // One critical section from submit to completion, so no other statement on the + // connection can end the query in between + ConnectionLockGuard conn_lock(*this); + auto res = PrepareAndSubmitInternal(std::move(last_statement), std::move(params)); if (!res) { return nullptr; } + { + D_ASSERT(duckdb::PyUtil::GilCheck()); + nb::gil_scoped_release release; + CompleteQuery(*res); + } if (res->GetStatementProperties().return_type != StatementReturnType::QUERY_RESULT) { return nullptr; } - if (res->GetResultType() == QueryResultType::STREAM_RESULT) { - auto &stream_result = res->Cast(); - res = stream_result.Materialize(); - } - auto &materialized_result = res->Cast(); - relation = make_shared_ptr(connection.context, materialized_result.TakeCollection(), - res->GetNames(), Identifier(alias)); + relation = make_shared_ptr(connection.context, res->TakeCollection(), res->GetNames(), + Identifier(alias)); } return CreateRelation(std::move(relation)); } diff --git a/src/pyrelation.cpp b/src/pyrelation.cpp index 24dda20e..38ff3984 100644 --- a/src/pyrelation.cpp +++ b/src/pyrelation.cpp @@ -12,10 +12,11 @@ #include "duckdb/parser/statement/pragma_statement.hpp" #include "duckdb/common/box_renderer.hpp" #include "duckdb/main/query_result.hpp" -#include "duckdb/main/materialized_query_result.hpp" #include "duckdb/parser/statement/explain_statement.hpp" #include "duckdb/catalog/default/default_types.hpp" #include "duckdb/main/relation/value_relation.hpp" +#include "duckdb/main/relation/table_function_relation.hpp" +#include "duckdb_python/map.hpp" #include "duckdb_python/expression/pyexpression.hpp" #include "duckdb/common/arrow/physical_arrow_collector.hpp" #include "duckdb_python/arrow/arrow_export_utils.hpp" @@ -262,7 +263,7 @@ void DuckDBPyRelation::AssertRelation() const { } void DuckDBPyRelation::AssertResultOpen() const { - if (!result || result->IsClosed()) { + if (!result) { throw InvalidInputException("No open result set"); } } @@ -807,33 +808,49 @@ duckdb::pyarrow::RecordBatchReader DuckDBPyRelation::FetchRecordBatchReader(idx_ return result->FetchRecordBatchReader(rows_per_batch); } -static unique_ptr PyExecuteRelation(const shared_ptr &rel, bool stream_result = false) { +//! Submits the relation and returns the undriven handle. Call with the GIL released. +static unique_ptr PySubmitRelation(const shared_ptr &rel, bool stream_result) { + auto context = rel->context->GetContext(); + QueryParameters parameters; + // A stream can only be opened on a handle whose retention is still undecided at submission + parameters.result_eagerness = stream_result ? ResultEagerness::AUTO : ResultEagerness::FORCED; + auto result = context->Submit(rel, parameters); + if (result->HasError()) { + result->ThrowError(); + } + return result; +} + +static unique_ptr PyExecuteRelation(const shared_ptr &rel) { if (!rel) { return nullptr; } - auto context = rel->context->GetContext(); D_ASSERT(duckdb::PyUtil::GilCheck()); nb::gil_scoped_release release; - auto pending_query = context->PendingQuery(rel, stream_result); - return DuckDBPyConnection::CompletePendingQuery(*pending_query); + auto query_result = PySubmitRelation(rel, false); + DuckDBPyConnection::CompleteQuery(*query_result); + return query_result; } -unique_ptr DuckDBPyRelation::ExecuteInternal(bool stream_result) { +unique_ptr DuckDBPyRelation::ExecuteInternal() { this->executed = true; - return PyExecuteRelation(rel, stream_result); + return PyExecuteRelation(rel); } void DuckDBPyRelation::ExecuteOrThrow(bool stream_result) { nb::gil_scoped_acquire gil; result.reset(); - auto query_result = ExecuteInternal(stream_result); - if (!query_result) { + if (!rel) { throw InternalException("ExecuteOrThrow - no query available to execute"); } - if (query_result->HasError()) { - query_result->ThrowError(); + this->executed = true; + std::shared_ptr py_result; + { + nb::gil_scoped_release release; + auto submitted = PySubmitRelation(rel, stream_result); + py_result = std::make_shared(std::move(submitted), stream_result); } - result = std::make_unique(std::move(query_result)); + result = std::move(py_result); } PandasDataFrame DuckDBPyRelation::FetchDF(bool date_as_object) { @@ -843,9 +860,6 @@ PandasDataFrame DuckDBPyRelation::FetchDF(bool date_as_object) { } ExecuteOrThrow(); } - if (result->IsClosed()) { - return nb::none(); - } auto df = result->FetchDF(date_as_object); result = nullptr; return df; @@ -858,9 +872,6 @@ Optional DuckDBPyRelation::FetchOne() { } ExecuteOrThrow(true); } - if (result->IsClosed()) { - return nb::none(); - } return result->Fetchone(); } @@ -872,9 +883,6 @@ nb::list DuckDBPyRelation::FetchMany(idx_t size) { ExecuteOrThrow(true); D_ASSERT(result); } - if (result->IsClosed()) { - return nb::list(); - } return result->Fetchmany(size); } @@ -885,9 +893,6 @@ nb::list DuckDBPyRelation::FetchAll() { } ExecuteOrThrow(); } - if (result->IsClosed()) { - return nb::list(); - } auto res = result->Fetchall(); result = nullptr; return res; @@ -900,9 +905,6 @@ nb::dict DuckDBPyRelation::FetchNumpy() { } ExecuteOrThrow(); } - if (result->IsClosed()) { - return nb::borrow(nb::none()); - } auto res = result->FetchNumpy(); result = nullptr; return res; @@ -915,9 +917,6 @@ nb::dict DuckDBPyRelation::FetchPyTorch() { } ExecuteOrThrow(); } - if (result->IsClosed()) { - return nb::borrow(nb::none()); - } auto res = result->FetchPyTorch(); result = nullptr; return res; @@ -930,15 +929,12 @@ nb::dict DuckDBPyRelation::FetchTF() { } ExecuteOrThrow(); } - if (result->IsClosed()) { - return nb::borrow(nb::none()); - } auto res = result->FetchTF(); result = nullptr; return res; } -nb::dict DuckDBPyRelation::FetchNumpyInternal(bool stream, idx_t vectors_per_chunk) { +nb::dict DuckDBPyRelation::FetchNumpyInternal(bool chunked, idx_t vectors_per_chunk) { if (!result) { if (!rel) { return nb::borrow(nb::none()); @@ -946,12 +942,11 @@ nb::dict DuckDBPyRelation::FetchNumpyInternal(bool stream, idx_t vectors_per_chu ExecuteOrThrow(); } AssertResultOpen(); - auto res = result->FetchNumpyInternal(stream, vectors_per_chunk); + auto res = result->FetchNumpyInternal(chunked, vectors_per_chunk); result = nullptr; return res; } -//! Should this also keep track of when the result is empty and set result->result_closed accordingly? PandasDataFrame DuckDBPyRelation::FetchDFChunk(idx_t vectors_per_chunk, bool date_as_object) { if (!result) { if (!rel) { @@ -994,8 +989,6 @@ nb::object DuckDBPyRelation::ToArrowCapsule(const nb::object &requested_schema) if (!rel) { return nb::none(); } - // 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(); @@ -1042,7 +1035,6 @@ duckdb::pyarrow::RecordBatchReader DuckDBPyRelation::ToRecordBatch(idx_t batch_s if (!rel) { return nb::none(); } - // Fresh relation: stream lazily on the user's own context (survives `del conn`). ExecuteOrThrow(true); } AssertResultOpen(); @@ -1562,7 +1554,7 @@ std::unique_ptr DuckDBPyRelation::Query(const string &view_nam { D_ASSERT(duckdb::PyUtil::GilCheck()); nb::gil_scoped_release release; - auto query_result = rel->context->GetContext()->Query(std::move(parser.statements[0]), false); + auto query_result = rel->context->GetContext()->Query(std::move(parser.statements[0]), QueryParameters()); // Execute it anyways, for creation/altering statements // We only care that it succeeds, we can't store the result D_ASSERT(query_result); @@ -1649,15 +1641,10 @@ void DuckDBPyRelation::Create(const string &table) { std::unique_ptr DuckDBPyRelation::Map(nb::callable fun, Optional schema) { AssertRelation(); - vector params; - params.emplace_back(Value::POINTER(CastPointerToValue(fun.ptr()))); - params.emplace_back(Value::POINTER(CastPointerToValue(schema.ptr()))); - auto relation = DeriveRelation(rel->TableFunction("python_map_function", params)); - auto rel_dependency = make_uniq(); - rel_dependency->AddDependency("map", PythonDependencyItem::Create(std::move(fun))); - rel_dependency->AddDependency("schema", PythonDependencyItem::Create(std::move(schema))); - relation->rel->AddExternalDependency(std::move(rel_dependency)); - return relation; + auto info = make_shared_ptr(std::move(fun), std::move(schema)); + return DeriveRelation(make_shared_ptr(rel->context->GetContext(), "python_map_function", + vector(), named_parameter_map_t(), rel, true, + std::move(info))); } string DuckDBPyRelation::ToStringInternal(const BoxRendererConfig &config, bool invalidate_cache) { @@ -1748,9 +1735,10 @@ string DuckDBPyRelation::Explain(ExplainType type, const string &format) { const bool auto_format = format.empty(); auto explain_format = auto_format ? GetExplainFormat(type) : ProfilerPrintFormat(format); auto res = rel->Explain(type, explain_format); - D_ASSERT(res->GetResultType() == duckdb::QueryResultType::MATERIALIZED_RESULT); - auto &materialized = res->Cast(); - auto &coll = materialized.Collection(); + if (res->HasError()) { + res->ThrowError(); + } + auto &coll = res->Collection(); // Only the implicit Jupyter path renders HTML inline; an explicitly requested format always returns a string. const bool jupyter_html = auto_format && explain_format == ProfilerPrintFormat::HTML() && DuckDBPyConnection::IsJupyter(); @@ -1770,7 +1758,7 @@ string DuckDBPyRelation::Explain(ExplainType type, const string &format) { return result_; } - auto chunk = materialized.Fetch(); + auto chunk = res->Fetch(); for (idx_t i = 0; i < chunk->size(); i++) { auto plan = chunk->GetValue(1, i); auto plan_string = plan.GetValue(); diff --git a/src/pyresult.cpp b/src/pyresult.cpp index 7afc8be4..f046e35c 100644 --- a/src/pyresult.cpp +++ b/src/pyresult.cpp @@ -6,20 +6,17 @@ #include "duckdb_python/arrow/arrow_array_stream.hpp" #include "duckdb/common/arrow/arrow.hpp" -#include "duckdb/common/arrow/arrow_util.hpp" #include "duckdb/common/arrow/arrow_converter.hpp" #include "duckdb/common/arrow/arrow_wrapper.hpp" #include "duckdb/common/arrow/result_arrow_wrapper.hpp" #include "duckdb/common/types/uuid.hpp" #include "duckdb/common/exception.hpp" -#include "duckdb/common/enums/stream_execution_result.hpp" #include "duckdb_python/arrow/arrow_export_utils.hpp" #include "duckdb/common/arrow/arrow_query_result.hpp" #include "duckdb/common/arrow/physical_arrow_collector.hpp" -#include "duckdb/main/chunk_scan_state/query_result.hpp" #include "duckdb/main/client_config.hpp" -#include "duckdb/main/materialized_query_result.hpp" -#include "duckdb/main/stream_query_result.hpp" +#include "duckdb/main/client_context.hpp" +#include "duckdb/main/query_result.hpp" #include "duckdb/parser/expression/star_expression.hpp" #include "duckdb/parser/query_node/select_node.hpp" #include "duckdb/parser/statement/select_statement.hpp" @@ -29,14 +26,67 @@ using namespace nanobind::literals; namespace duckdb { -DuckDBPyResult::DuckDBPyResult(unique_ptr result_p) : result(std::move(result_p)) { +DuckDBPyResult::DuckDBPyResult(unique_ptr completed) : result(std::move(completed)) { if (!result) { throw InternalException("PyResult created without a result object"); } } +DuckDBPyResult::DuckDBPyResult(unique_ptr submitted, bool stream_result) { + if (!submitted) { + throw InternalException("PyResult created without a result object"); + } + if (submitted->HasError()) { + submitted->ThrowError(); + } + // A statement the planner settles on retained cannot be drained, and neither can a result a + // delegating collector already built, which carries no buffer. + const bool can_stream = + submitted->HasBufferedData() && submitted->GetStatementProperties().result_eagerness != ResultEagerness::FORCED; + if (stream_result && can_stream) { + this->submitted = std::move(submitted); + return; + } + DuckDBPyConnection::CompleteQuery(*submitted); + result = std::move(submitted); +} + +void DuckDBPyResult::EnsureStream() { + if (submitted) { + stream = make_uniq(std::move(submitted)); + } +} + const vector &DuckDBPyResult::ResultNames() const { - return names_override.empty() ? result->GetNames() : names_override; + if (!names_override.empty()) { + return names_override; + } + if (stream) { + return stream->GetNames(); + } + return submitted ? submitted->GetNames() : result->GetNames(); +} + +void DuckDBPyResult::CloseStream() { + if (!stream && !submitted) { + return; + } + auto close = [&]() { + if (stream) { + stream->Close(); + } + if (submitted) { + submitted->Close(); + } + }; + // Ending the query waits for its running tasks, and a task inside a Python UDF cannot finish + // until the GIL is free. + if (duckdb::PyUtil::GilCheck()) { + nb::gil_scoped_release release; + close(); + } else { + close(); + } } DuckDBPyResult::~DuckDBPyResult() { @@ -48,99 +98,160 @@ DuckDBPyResult::~DuckDBPyResult() { // to run without a valid PyThreadState — see duckdb-python#456. try { D_ASSERT(duckdb::PyUtil::GilCheck()); + CloseStream(); + stream.reset(); + submitted.reset(); result.reset(); current_chunk.reset(); } catch (...) { // NOLINT } } -ClientProperties DuckDBPyResult::GetClientProperties() { - return result->client_properties; +const ClientProperties &DuckDBPyResult::GetClientProperties() const { + if (stream) { + return stream->GetClientProperties(); + } + return submitted ? submitted->client_properties : result->client_properties; } vector DuckDBPyResult::GetNames() { - if (!result) { + if (Empty()) { throw InternalException("Calling GetNames without a result object"); } return IdentifiersToStrings(ResultNames()); } -const vector &DuckDBPyResult::GetTypes() { - if (!result) { +const vector &DuckDBPyResult::GetTypes() const { + if (Empty()) { throw InternalException("Calling GetTypes without a result object"); } - return result->GetTypes(); + if (stream) { + return stream->GetTypes(); + } + return submitted ? submitted->GetTypes() : result->GetTypes(); } unique_ptr DuckDBPyResult::FetchChunk() { - if (!result) { + if (Empty()) { throw InternalException("FetchChunk called without a result object"); } - return FetchNext(*result); + return FetchNext(); } -unique_ptr DuckDBPyResult::FetchNext(QueryResult &query_result) { - if (!result_closed && query_result.GetResultType() == QueryResultType::STREAM_RESULT && - !query_result.Cast().IsOpen()) { - result_closed = true; - return nullptr; - } - if (query_result.GetResultType() == QueryResultType::STREAM_RESULT) { - auto &stream_result = query_result.Cast(); - StreamExecutionResult execution_result; - while (!StreamQueryResult::IsChunkReady(execution_result = stream_result.ExecuteTask())) { - { - nb::gil_scoped_acquire gil; - if (PyErr_CheckSignals() != 0) { - throw std::runtime_error("Query interrupted"); - } - } - if (execution_result == StreamExecutionResult::BLOCKED) { - stream_result.WaitForTask(); +unique_ptr DuckDBPyResult::FetchStreamChunk() { + while (true) { + unique_ptr chunk; + auto state = stream->TryFetch(chunk); + if (chunk) { + return chunk; + } + if (state == QueryResultState::FINISHED) { + return nullptr; + } + if (state == QueryResultState::EXECUTION_ERROR) { + stream->GetErrorObject().Throw(); + } + { + nb::gil_scoped_acquire gil; + if (PyErr_CheckSignals() != 0) { + throw std::runtime_error("Query interrupted"); } } - if (execution_result == StreamExecutionResult::EXECUTION_CANCELLED) { - throw InvalidInputException("The execution of the query was cancelled before it could finish, likely " - "caused by executing a different query"); + state = stream->ExecuteTask(); + if (state == QueryResultState::BLOCKED || state == QueryResultState::NO_TASKS_AVAILABLE) { + stream->WaitForTask(); } - if (execution_result == StreamExecutionResult::EXECUTION_ERROR) { - stream_result.ThrowError(); + } +} + +unique_ptr DuckDBPyResult::FetchNext() { + EnsureStream(); + if (stream) { + auto chunk = FetchStreamChunk(); + if (chunk) { + chunk->Flatten(); } + return chunk; } - auto chunk = query_result.Fetch(); - if (query_result.HasError()) { - query_result.ThrowError(); + auto chunk = result->Fetch(); + if (result->HasError()) { + result->ThrowError(); } return chunk; } -unique_ptr DuckDBPyResult::FetchNextRaw(QueryResult &query_result) { - if (!result_closed && query_result.GetResultType() == QueryResultType::STREAM_RESULT && - !query_result.Cast().IsOpen()) { - result_closed = true; - return nullptr; +unique_ptr DuckDBPyResult::FetchNextRaw() { + EnsureStream(); + if (stream) { + return FetchStreamChunk(); } - auto chunk = query_result.FetchRaw(); - if (query_result.HasError()) { - query_result.ThrowError(); + auto chunk = result->FetchRaw(); + if (result->HasError()) { + result->ThrowError(); } return chunk; } +unique_ptr DuckDBPyResult::TakeBufferedRows() { + unique_ptr remainder; + if (current_chunk && chunk_offset < current_chunk->size() && !StreamEnded()) { + remainder = make_uniq(); + remainder->Initialize(Allocator::DefaultAllocator(), current_chunk->GetTypes()); + current_chunk->Copy(*remainder, chunk_offset); + } + current_chunk.reset(); + chunk_offset = 0; + return remainder; +} + +void DuckDBPyResult::Retain() { + if (submitted) { + { + D_ASSERT(duckdb::PyUtil::GilCheck()); + nb::gil_scoped_release release; + DuckDBPyConnection::CompleteQuery(*submitted); + } + result = std::move(submitted); + return; + } + if (!stream) { + return; + } + auto collection = make_uniq(Allocator::DefaultAllocator(), stream->GetTypes()); + { + D_ASSERT(duckdb::PyUtil::GilCheck()); + nb::gil_scoped_release release; + if (auto buffered = TakeBufferedRows()) { + collection->Append(*buffered); + } + while (auto chunk = FetchStreamChunk()) { + collection->Append(*chunk); + } + } + auto retained = make_uniq(stream->GetStatementType(), stream->GetStatementProperties(), + stream->GetNames(), std::move(collection), stream->GetClientProperties()); + CloseStream(); + stream.reset(); + result = std::move(retained); +} + Optional DuckDBPyResult::Fetchone() { - if (!result) { + if (Empty()) { throw InvalidInputException("result closed"); } - if (!current_chunk || chunk_offset >= current_chunk->size()) { + { nb::gil_scoped_release release; - current_chunk = FetchNext(*result); - chunk_offset = 0; + if (!current_chunk || chunk_offset >= current_chunk->size() || StreamEnded()) { + current_chunk = FetchNext(); + chunk_offset = 0; + } } if (!current_chunk || current_chunk->size() == 0) { return nb::none(); } - auto &types = result->GetTypes(); + auto &types = GetTypes(); + auto &client_properties = GetClientProperties(); duckdb::PyUtil::TupleBuilder row(types.size()); for (idx_t col_idx = 0; col_idx < types.size(); col_idx++) { auto &mask = FlatVector::Validity(current_chunk->data[col_idx]); @@ -148,7 +259,7 @@ Optional DuckDBPyResult::Fetchone() { row.append(nb::none()); } else { auto val = current_chunk->data[col_idx].GetValue(chunk_offset); - row.append(PythonObject::FromValue(val, types[col_idx], result->client_properties)); + row.append(PythonObject::FromValue(val, types[col_idx], client_properties)); } } chunk_offset++; @@ -184,7 +295,7 @@ nb::dict DuckDBPyResult::FetchNumpy() { } void DuckDBPyResult::FillNumpy(nb::dict &res, idx_t col_idx, NumpyResultConversion &conversion, const char *name) { - if (result->GetTypes()[col_idx].id() == LogicalTypeId::ENUM) { + if (GetTypes()[col_idx].id() == LogicalTypeId::ENUM) { auto &import_cache = *DuckDBPyConnection::ImportCache(); auto pandas_categorical = import_cache.pandas.Categorical(); auto categorical_dtype = import_cache.pandas.CategoricalDtype(); @@ -208,8 +319,7 @@ void DuckDBPyResult::FillNumpy(nb::dict &res, idx_t col_idx, NumpyResultConversi } } -void InsertCategory(QueryResult &result, unordered_map &categories) { - auto &types = result.GetTypes(); +void InsertCategory(const vector &types, unordered_map &categories) { for (idx_t col_idx = 0; col_idx < types.size(); col_idx++) { auto &type = types[col_idx]; if (type.id() == LogicalTypeId::ENUM) { @@ -226,63 +336,59 @@ void InsertCategory(QueryResult &result, unordered_map &categor } std::unique_ptr DuckDBPyResult::InitializeNumpyConversion(bool pandas) { - if (!result) { + if (Empty()) { throw InvalidInputException("result closed"); } idx_t initial_capacity = STANDARD_VECTOR_SIZE * 2ULL; - if (result->GetResultType() == QueryResultType::MATERIALIZED_RESULT) { - // materialized query result: we know exactly how much space we need - auto &materialized = result->Cast(); - initial_capacity = materialized.RowCount(); + if (result && result->GetResultType() == QueryResultType::MATERIALIZED_RESULT) { + initial_capacity = result->RowCount(); } - auto conversion = std::make_unique(result->GetTypes(), initial_capacity, - result->client_properties, pandas); + auto conversion = + std::make_unique(GetTypes(), initial_capacity, GetClientProperties(), pandas); return conversion; } -nb::dict DuckDBPyResult::FetchNumpyInternal(bool stream, idx_t vectors_per_chunk, +nb::dict DuckDBPyResult::FetchNumpyInternal(bool chunked, idx_t vectors_per_chunk, std::unique_ptr conversion_p) { - if (!result) { + if (Empty()) { throw InvalidInputException("result closed"); } if (!conversion_p) { conversion_p = InitializeNumpyConversion(); } auto &conversion = *conversion_p; + if (!chunked) { + vectors_per_chunk = NumericLimits::Maximum(); + } - if (result->GetResultType() == QueryResultType::MATERIALIZED_RESULT) { - auto &materialized = result->Cast(); - for (auto &chunk : materialized.Collection().Chunks()) { - conversion.Append(chunk); + idx_t count_vec = 0; + if (vectors_per_chunk > 0) { + unique_ptr buffered; + { + D_ASSERT(duckdb::PyUtil::GilCheck()); + nb::gil_scoped_release release; + buffered = TakeBufferedRows(); } - InsertCategory(materialized, categories); - materialized.Collection().Reset(); - } else { - D_ASSERT(result->GetResultType() == QueryResultType::STREAM_RESULT); - if (!stream) { - vectors_per_chunk = NumericLimits::Maximum(); + if (buffered) { + conversion.Append(*buffered); + count_vec++; } - auto &stream_result = result->Cast(); - for (idx_t count_vec = 0; count_vec < vectors_per_chunk; count_vec++) { - if (!stream_result.IsOpen()) { - break; - } - unique_ptr chunk; - { - D_ASSERT(duckdb::PyUtil::GilCheck()); - nb::gil_scoped_release release; - chunk = FetchNextRaw(stream_result); - } - if (!chunk || chunk->size() == 0) { - //! finished - break; - } - conversion.Append(*chunk); - InsertCategory(stream_result, categories); + } + for (; count_vec < vectors_per_chunk; count_vec++) { + unique_ptr chunk; + { + D_ASSERT(duckdb::PyUtil::GilCheck()); + nb::gil_scoped_release release; + chunk = FetchNextRaw(); + } + if (!chunk || chunk->size() == 0) { + break; } + conversion.Append(*chunk); } + InsertCategory(GetTypes(), categories); // now that we have materialized the result in contiguous arrays, construct the actual NumPy arrays or categorical // types @@ -305,14 +411,15 @@ static void ReplaceDFColumn(PandasDataFrame &df, const char *col_name, idx_t idx void DuckDBPyResult::ConvertDateTimeTypes(PandasDataFrame &df, bool date_as_object) const { auto names = nb::cast>(df.attr("columns")); - for (idx_t i = 0; i < result->ColumnCount(); i++) { - if (result->GetTypes()[i] == LogicalType::TIMESTAMP_TZ) { + auto &types = GetTypes(); + for (idx_t i = 0; i < types.size(); i++) { + if (types[i] == LogicalType::TIMESTAMP_TZ) { // first localize to UTC then convert to timezone_config auto utc_local = df[names[i].c_str()].attr("dt").attr("tz_localize")("UTC"); - auto new_value = utc_local.attr("dt").attr("tz_convert")(result->client_properties.time_zone); + auto new_value = utc_local.attr("dt").attr("tz_convert")(GetClientProperties().time_zone); // We need to create the column anew because the exact dt changed to a new timezone ReplaceDFColumn(df, names[i].c_str(), i, new_value); - } else if (date_as_object && result->GetTypes()[i] == LogicalType::DATE) { + } else if (date_as_object && types[i] == LogicalType::DATE) { nb::object new_value = df[names[i].c_str()].attr("dt").attr("date"); ReplaceDFColumn(df, names[i].c_str(), i, new_value); } @@ -394,7 +501,7 @@ PandasDataFrame DuckDBPyResult::FrameFromNumpy(bool date_as_object, const nb::ha ConvertDateTimeTypes(df, date_as_object); auto names = nb::cast>(df.attr("columns")); - D_ASSERT(result->ColumnCount() == names.size()); + D_ASSERT(GetTypes().size() == names.size()); return df; } @@ -426,10 +533,8 @@ nb::dict DuckDBPyResult::FetchTF() { return result_dict; } -// `SELECT * FROM ` over `collection`, executed as a SelectStatement rather -// than via PendingQuery(relation) - the latter's RelationStatement stringifies the whole -// collection (O(rows)). The ColumnDataRef owns the collection, so it outlives the result -// (needed by the lazy stream path). +// A SelectStatement over a ColumnDataRef rather than a relation: a RelationStatement would +// stringify the whole collection when the query is submitted. static unique_ptr MakeColumnDataScanStatement(unique_ptr collection, const vector &names) { // The binder rejects duplicate column names; callers restore the originals afterwards. @@ -445,8 +550,6 @@ static unique_ptr MakeColumnDataScanStatement(unique_ptr ArrowQueryResult. void DuckDBPyResult::PromoteMaterializedToArrow(idx_t batch_size) { D_ASSERT(result->GetResultType() == QueryResultType::MATERIALIZED_RESULT); auto client_context = result->client_properties.client_context; @@ -454,9 +557,8 @@ void DuckDBPyResult::PromoteMaterializedToArrow(idx_t batch_size) { throw InternalException("Cannot promote result to Arrow: the originating client context is gone"); } auto context = client_context->shared_from_this(); - auto &materialized = result->Cast(); auto names = ResultNames(); - auto select = MakeColumnDataScanStatement(materialized.TakeCollection(), names); + auto select = MakeColumnDataScanStatement(result->TakeCollection(), names); auto &config = ClientConfig::GetConfig(*context); ScopedConfigSetting scoped_setting( @@ -472,11 +574,8 @@ void DuckDBPyResult::PromoteMaterializedToArrow(idx_t batch_size) { { D_ASSERT(duckdb::PyUtil::GilCheck()); nb::gil_scoped_release release; - auto pending_query = context->PendingQuery(std::move(select), QueryParameters(false)); - new_result = DuckDBPyConnection::CompletePendingQuery(*pending_query); - } - if (new_result->HasError()) { - new_result->ThrowError(); + new_result = context->Submit(std::move(select), QueryParameters()); + DuckDBPyConnection::CompleteQuery(*new_result); } names_override = std::move(names); // restore names de-duplicated by re-binding result = std::move(new_result); @@ -484,11 +583,12 @@ void DuckDBPyResult::PromoteMaterializedToArrow(idx_t batch_size) { template T DuckDBPyResult::RunWithArrowSchema(const std::function &fun, bool dedup_col_names) { - D_ASSERT(result); - if (!result->client_properties.client_context) { + D_ASSERT(!Empty()); + auto client_properties = GetClientProperties(); + if (!client_properties.client_context) { throw ConnectionException("Cannot fetch arrow schema without a valid connection"); } - auto ctx = result->client_properties.client_context->shared_from_this(); + auto ctx = client_properties.client_context->shared_from_this(); auto identifiers = ResultNames(); if (dedup_col_names) { @@ -498,13 +598,14 @@ T DuckDBPyResult::RunWithArrowSchema(const std::function ArrowSchema arrow_schema; ctx->RunFunctionInTransaction( - [&] { ArrowConverter::ToArrowSchema(&arrow_schema, result->GetTypes(), names, result->client_properties); }); + [&] { ArrowConverter::ToArrowSchema(&arrow_schema, GetTypes(), names, client_properties); }); return fun(arrow_schema); } duckdb::pyarrow::Table DuckDBPyResult::MaterializedResultToArrowTable(const ArrowSchema &arrow_schema, const idx_t rows_per_batch) { + Retain(); D_ASSERT(result); D_ASSERT(result->GetResultType() == QueryResultType::MATERIALIZED_RESULT || result->GetResultType() == QueryResultType::ARROW_RESULT); @@ -525,118 +626,172 @@ duckdb::pyarrow::Table DuckDBPyResult::MaterializedResultToArrowTable(const Arro } duckdb::pyarrow::Table DuckDBPyResult::FetchArrowTable(const idx_t rows_per_batch, const bool to_polars) { - if (!result) { + if (Empty()) { throw InvalidInputException("There is no query result"); } return RunWithArrowSchema( [&](const ArrowSchema &schema) -> duckdb::pyarrow::Table { - if (result->GetResultType() == QueryResultType::MATERIALIZED_RESULT || - result->GetResultType() == QueryResultType::ARROW_RESULT) { - return MaterializedResultToArrowTable(schema, rows_per_batch); - } - if (result->GetResultType() != QueryResultType::STREAM_RESULT) { - throw InternalException("FetchArrowTable called with unsupported query result: %d", - result->GetResultType()); - } - auto pyarrow_schema = pyarrow::ToPyArrowSchema(schema); - nb::list batches; - QueryResultChunkScanState scan_state(*result); - while (true) { - ArrowArray data; - idx_t count; - { - D_ASSERT(duckdb::PyUtil::GilCheck()); - nb::gil_scoped_release release; - count = ArrowUtil::FetchChunk(scan_state, result->client_properties, rows_per_batch, &data, - ArrowTypeExtensionData::GetExtensionTypes( - *result->client_properties.client_context, result->GetTypes())); - } - if (count == 0) { - break; - } - TransformDuckToArrowChunk(pyarrow_schema, data, batches); - } - return pyarrow::ToArrowTable(std::move(batches), pyarrow_schema); + return MaterializedResultToArrowTable(schema, rows_per_batch); }, to_polars); } -ArrowArrayStream DuckDBPyResult::FetchArrowArrayStream(idx_t rows_per_batch) { - if (!result) { - throw InvalidInputException("There is no query result"); +static void CheckBatchSize(idx_t rows_per_batch) { + if (rows_per_batch == 0) { + throw std::runtime_error("Approximate Batch Size of Record Batch MUST be higher than 0"); + } +} + +namespace { + +//! pyarrow releases an imported stream with the GIL held, and releasing the engine's stream ends +//! the query and joins its tasks, one of which may be inside a Python UDF waiting for the GIL. +//! Pulling batches has the same shape when a caller holds the GIL, so every callback drops it. +template +auto WithoutGil(FUN &&fun) -> decltype(fun()) { + if (duckdb::PyUtil::GilCheck()) { + nb::gil_scoped_release release; + return fun(); + } + return fun(); +} + +//! The engine's stream and, for a result whose query already ended, the context its batches are +//! converted under. An open stream keeps its own context; a retained result released it. +struct EngineStreamHandoff { + ArrowArrayStream engine; + shared_ptr context; +}; + +ArrowArrayStream &EngineStream(ArrowArrayStream *stream) { + return static_cast(stream->private_data)->engine; +} + +int ForwardGetSchema(ArrowArrayStream *stream, ArrowSchema *out) { + auto &engine = EngineStream(stream); + return WithoutGil([&]() { return engine.get_schema(&engine, out); }); +} + +int ForwardGetNext(ArrowArrayStream *stream, ArrowArray *out) { + auto &engine = EngineStream(stream); + return WithoutGil([&]() { return engine.get_next(&engine, out); }); +} + +const char *ForwardGetLastError(ArrowArrayStream *stream) { + auto &engine = EngineStream(stream); + return engine.get_last_error(&engine); +} + +void ForwardRelease(ArrowArrayStream *stream) { + if (!stream->release) { + return; } - if (result->GetResultType() != QueryResultType::STREAM_RESULT) { - throw InternalException("FetchArrowArrayStream called with unsupported query result: %d", - result->GetResultType()); + auto handoff = static_cast(stream->private_data); + WithoutGil([&]() { + if (handoff->engine.release) { + handoff->engine.release(&handoff->engine); + } + delete handoff; + }); + stream->private_data = nullptr; + stream->release = nullptr; +} + +//! Releases a stream that was never handed over, for example when the import into pyarrow throws +struct ArrowArrayStreamGuard { + ArrowArrayStream stream; + ~ArrowArrayStreamGuard() { + if (stream.release) { + stream.release(&stream); + } } - // The wrapper is owned by the ArrowArrayStream's private_data (released with the stream). - const auto result_stream = new ResultArrowArrayStreamWrapper(std::move(result), rows_per_batch); - return result_stream->stream; +}; + +} // namespace + +ArrowArrayStream DuckDBPyResult::FetchArrowArrayStream(idx_t rows_per_batch) { + if (stream) { + Retain(); + } + auto &client_context = GetClientProperties().client_context; + auto context = client_context ? client_context->shared_from_this() : shared_ptr(); + auto handle = submitted ? std::move(submitted) : std::move(result); + current_chunk.reset(); + chunk_offset = 0; + const auto result_stream = new ResultArrowArrayStreamWrapper(std::move(handle), rows_per_batch); + auto handoff = new EngineStreamHandoff {result_stream->stream, std::move(context)}; + ArrowArrayStream forwarding; + forwarding.get_schema = ForwardGetSchema; + forwarding.get_next = ForwardGetNext; + forwarding.get_last_error = ForwardGetLastError; + forwarding.release = ForwardRelease; + forwarding.private_data = handoff; + return forwarding; +} + +//! An Arrow result was converted already and has no collection the engine's stream could read +bool DuckDBPyResult::IsArrow() const { + return result && result->GetResultType() == QueryResultType::ARROW_RESULT; } duckdb::pyarrow::RecordBatchReader DuckDBPyResult::FetchRecordBatchReader(idx_t rows_per_batch) { - if (!result) { + if (Empty()) { throw InvalidInputException("There is no query result"); } + CheckBatchSize(rows_per_batch); - if (result->GetResultType() == QueryResultType::MATERIALIZED_RESULT || - result->GetResultType() == QueryResultType::ARROW_RESULT) { + if (IsArrow()) { constexpr bool dedup_column_names = false; - return RunWithArrowSchema( + auto reader = RunWithArrowSchema( [&](const ArrowSchema &schema) -> duckdb::pyarrow::RecordBatchReader { const auto table = MaterializedResultToArrowTable(schema, rows_per_batch); return nb::cast( table.attr("to_reader")(nb::arg("max_chunksize") = rows_per_batch)); }, dedup_column_names); + result.reset(); + return reader; } - if (result->GetResultType() != QueryResultType::STREAM_RESULT) { - throw InternalException("FetchRecordBatchReader called with unsupported query result: %d", - result->GetResultType()); - } - nb::gil_scoped_acquire acquire; auto pyarrow_lib_module = nb::module_::import_("pyarrow").attr("lib"); auto record_batch_reader_func = pyarrow_lib_module.attr("RecordBatchReader").attr("_import_from_c"); - auto stream = FetchArrowArrayStream(rows_per_batch); - nb::object record_batch_reader = record_batch_reader_func((uint64_t)&stream); // NOLINT + ArrowArrayStreamGuard guard {FetchArrowArrayStream(rows_per_batch)}; + nb::object record_batch_reader = record_batch_reader_func((uint64_t)&guard.stream); // NOLINT return nb::cast(record_batch_reader); } static void ArrowArrayStreamPyCapsuleDestructor(void *data) noexcept { - // nanobind capsule cleanup receives the raw pointer (via PyCapsule_GetPointer using the capsule's name) if (!data) { return; } - auto stream = reinterpret_cast(data); - if (stream->release) { - stream->release(stream); + auto arrow_stream = reinterpret_cast(data); + if (arrow_stream->release) { + arrow_stream->release(arrow_stream); } - delete stream; + delete arrow_stream; } nb::object DuckDBPyResult::FetchArrowCapsule(const idx_t rows_per_batch) { - if (!result) { + if (Empty()) { throw InvalidInputException("There is no query result"); } + CheckBatchSize(rows_per_batch); - constexpr bool dedup_column_names = false; - if (result->GetResultType() == QueryResultType::MATERIALIZED_RESULT || - result->GetResultType() == QueryResultType::ARROW_RESULT) { - return RunWithArrowSchema( + if (IsArrow()) { + constexpr bool dedup_column_names = false; + auto capsule = RunWithArrowSchema( [&](const ArrowSchema &schema) -> nb::object { const auto table = MaterializedResultToArrowTable(schema, rows_per_batch); return table.attr("__arrow_c_stream__")(); }, dedup_column_names); - } - if (result->GetResultType() != QueryResultType::STREAM_RESULT) { - throw InternalException("FetchArrowCapsule called with unsupported query result: %d", result->GetResultType()); + result.reset(); + return capsule; } auto inner_stream = FetchArrowArrayStream(rows_per_batch); - auto stream = new ArrowArrayStream(); - *stream = inner_stream; - return nb::capsule(stream, "arrow_array_stream", ArrowArrayStreamPyCapsuleDestructor); + auto arrow_stream = new ArrowArrayStream(); + *arrow_stream = inner_stream; + return nb::capsule(arrow_stream, "arrow_array_stream", ArrowArrayStreamPyCapsuleDestructor); } nb::list DuckDBPyResult::GetDescription(const vector &names, const vector &types) { @@ -651,11 +806,12 @@ nb::list DuckDBPyResult::GetDescription(const vector &names, const vecto } void DuckDBPyResult::Close() { - result = nullptr; -} - -bool DuckDBPyResult::IsClosed() const { - return result_closed; + CloseStream(); + stream.reset(); + submitted.reset(); + result.reset(); + current_chunk.reset(); + chunk_offset = 0; } } // namespace duckdb diff --git a/src/python_dependency.cpp b/src/python_dependency.cpp index 2b2f82e0..46566458 100644 --- a/src/python_dependency.cpp +++ b/src/python_dependency.cpp @@ -6,11 +6,6 @@ namespace duckdb { PythonDependencyItem::PythonDependencyItem(unique_ptr &&object) : object(std::move(object)) { } -PythonDependencyItem::~PythonDependencyItem() { // NOLINT - cannot throw in exception - nb::gil_scoped_acquire gil; - object.reset(); -} - shared_ptr PythonDependencyItem::Create(nb::object object) { auto registered_object = make_uniq(std::move(object)); return make_shared_ptr(std::move(registered_object)); diff --git a/src/python_replacement_scan.cpp b/src/python_replacement_scan.cpp index 4005dbd4..f205d870 100644 --- a/src/python_replacement_scan.cpp +++ b/src/python_replacement_scan.cpp @@ -53,14 +53,7 @@ static void CreateArrowScan(const string &name, nb::object entry, TableFunctionR auto dependency_item = PythonDependencyItem::Create(stream_messages); external_dependency->AddDependency("replacement_cache", std::move(dependency_item)); } else { - auto stream_factory = make_uniq(entry.ptr(), client_properties, type); - auto stream_factory_produce = PythonTableArrowArrayStreamFactory::Produce; - auto stream_factory_get_schema = PythonTableArrowArrayStreamFactory::GetSchema; - - children.push_back(ConstantExpression::FromValue(Value::POINTER(CastPointerToValue(stream_factory.get())))); - children.push_back(ConstantExpression::FromValue(Value::POINTER(CastPointerToValue(stream_factory_produce)))); - children.push_back( - ConstantExpression::FromValue(Value::POINTER(CastPointerToValue(stream_factory_get_schema)))); + table_function.bind_info = make_shared_ptr(entry, client_properties, type); if (type == PyArrowObjectType::PyCapsule) { // Disable projection+filter pushdown for bare capsules (single-use, no PyArrow wrapper) @@ -77,9 +70,7 @@ static void CreateArrowScan(const string &name, nb::object entry, TableFunctionR } else { table_function.function = make_uniq("arrow_scan", std::move(children)); } - auto dependency_item = - PythonDependencyItem::Create(make_uniq(std::move(stream_factory), entry)); - external_dependency->AddDependency("replacement_cache", std::move(dependency_item)); + external_dependency->AddDependency("replacement_cache", PythonDependencyItem::Create(entry)); } table_function.external_dependency = std::move(external_dependency); } @@ -124,11 +115,10 @@ unique_ptr PythonReplacementScan::TryReplacementObject(const nb::objec } else { string name = "df_" + StringUtil::GenerateRandomName(); auto new_df = PandasScanFunction::PandasReplaceCopiedNames(entry); - children.push_back(ConstantExpression::FromValue(Value::POINTER(CastPointerToValue(new_df.ptr())))); table_function->function = make_uniq("pandas_scan", std::move(children)); + table_function->bind_info = make_shared_ptr(std::move(new_df)); auto dependency = make_uniq(); dependency->AddDependency("replacement_cache", PythonDependencyItem::Create(entry)); - dependency->AddDependency("copy", PythonDependencyItem::Create(new_df)); table_function->external_dependency = std::move(dependency); } } else if (DuckDBPyRelation::IsRelation(entry)) { @@ -192,11 +182,10 @@ unique_ptr PythonReplacementScan::TryReplacementObject(const nb::objec throw NotImplementedException("Unsupported Numpy object"); break; } - children.push_back(ConstantExpression::FromValue(Value::POINTER(CastPointerToValue(data.ptr())))); table_function->function = make_uniq("pandas_scan", std::move(children)); + table_function->bind_info = make_shared_ptr(std::move(data)); auto dependency = make_uniq(); dependency->AddDependency("replacement_cache", PythonDependencyItem::Create(entry)); - dependency->AddDependency("data", PythonDependencyItem::Create(data)); table_function->external_dependency = std::move(dependency); } else { // This throws an error later on! diff --git a/src/python_udf.cpp b/src/python_udf.cpp index b6342ec2..3416b349 100644 --- a/src/python_udf.cpp +++ b/src/python_udf.cpp @@ -79,16 +79,13 @@ void AreExtensionsRegistered(const LogicalType &arrow_type, const LogicalType &d } } static void ConvertArrowTableToVector(const nb::object &table, Vector &out, ClientContext &context, idx_t count) { - // Create the stream factory from the Table object - auto ptr = table.ptr(); D_ASSERT(duckdb::PyUtil::GilCheck()); + // Declared before the GIL release so the factory's last reference drops with the GIL held + TableFunctionRef empty; + empty.bind_info = make_shared_ptr(table, context.GetClientProperties(), + PyArrowObjectType::Table); nb::gil_scoped_release gil; - auto stream_factory = - make_uniq(ptr, context.GetClientProperties(), PyArrowObjectType::Table); - auto stream_factory_produce = PythonTableArrowArrayStreamFactory::Produce; - auto stream_factory_get_schema = PythonTableArrowArrayStreamFactory::GetSchema; - // Get the functions we need auto function = ArrowTableFunction::ArrowScanFunction; auto bind = ArrowTableFunction::ArrowScanBind; @@ -97,15 +94,10 @@ static void ConvertArrowTableToVector(const nb::object &table, Vector &out, Clie // Prepare the inputs for the bind vector children; - children.reserve(3); - children.push_back(Value::POINTER(CastPointerToValue(stream_factory.get()))); - children.push_back(Value::POINTER(CastPointerToValue(stream_factory_produce))); - children.push_back(Value::POINTER(CastPointerToValue(stream_factory_get_schema))); named_parameter_map_t named_params; vector input_types; vector input_names; - TableFunctionRef empty; TableFunction dummy_table_function; dummy_table_function.name = "ConvertArrowTableToVector"; TableFunctionBindInput bind_input(children, named_params, input_types, input_names, nullptr, nullptr, diff --git a/tests/fast/api/test_query_interrupt.py b/tests/fast/api/test_query_interrupt.py index 4a5a02e5..db41a18e 100644 --- a/tests/fast/api/test_query_interrupt.py +++ b/tests/fast/api/test_query_interrupt.py @@ -8,6 +8,11 @@ import duckdb +def drain(res): + while res.fetchone() is not None: + pass + + def send_keyboard_interrupt(): # Wait a little, so we're sure the 'execute' has started time.sleep(0.1) @@ -34,3 +39,36 @@ def test_query_interruption(self): except KeyboardInterrupt: pytest.fail("Interrupted by user") thread.join() + + @pytest.mark.xfail( + condition=platform.system() == "Emscripten", + reason="Emscripten builds cannot use threads", + ) + @pytest.mark.timeout(120) + def test_streaming_fetch_interruption(self): + con = duckdb.connect() + con.execute("SET max_streaming_buffer_size='1MB'") + res = con.sql("select i from range(8000000000) t(i)") + assert res.fetchone() == (0,) + + interrupter = threading.Thread(target=send_keyboard_interrupt) + interrupter.start() + # The interrupt lands either in the fetch loop's own signal check or between two fetches in + # the interpreter, so either exception proves the drain stopped. + with pytest.raises((RuntimeError, KeyboardInterrupt)): + drain(res) + interrupter.join() + + @pytest.mark.xfail( + condition=platform.system() == "Emscripten", + reason="Emscripten builds cannot use threads", + ) + @pytest.mark.timeout(120) + def test_materializing_fetch_interruption(self): + con = duckdb.connect() + rel = con.sql("select count(*) from range(8000000000)") + interrupter = threading.Thread(target=send_keyboard_interrupt) + interrupter.start() + with pytest.raises(RuntimeError, match="Query interrupted"): + rel.df() + interrupter.join() diff --git a/tests/fast/api/test_streaming_result.py b/tests/fast/api/test_streaming_result.py index d42939bf..9bf8ece6 100644 --- a/tests/fast/api/test_streaming_result.py +++ b/tests/fast/api/test_streaming_result.py @@ -1,6 +1,9 @@ +import gc + import pytest import duckdb +from duckdb.sqltypes import BIGINT class TestStreamingResult: @@ -13,9 +16,9 @@ def test_fetch_one(self, duckdb_cursor): result.append(tpl[0]) assert result == list(range(5000)) - # fetch one with error + # fetch one with error: the bad row sits inside the first chunk res = duckdb_cursor.sql( - "SELECT CASE WHEN i < 10000 THEN i ELSE concat('hello', i::VARCHAR)::INT END FROM range(100000) t(i)" + "SELECT CASE WHEN i < 1000 THEN i ELSE concat('hello', i::VARCHAR)::INT END FROM range(100000) t(i)" ) with pytest.raises(duckdb.ConversionException): res.fetchone() @@ -29,9 +32,9 @@ def test_fetch_many(self, duckdb_cursor): result += [x[0] for x in tpl] assert result == list(range(5000)) - # fetch many with error + # fetch many with error: the bad row sits inside the first chunk res = duckdb_cursor.sql( - "SELECT CASE WHEN i < 10000 THEN i ELSE concat('hello', i::VARCHAR)::INT END FROM range(100000) t(i)" + "SELECT CASE WHEN i < 1000 THEN i ELSE concat('hello', i::VARCHAR)::INT END FROM range(100000) t(i)" ) with pytest.raises(duckdb.ConversionException): res.fetchmany(10) @@ -75,3 +78,188 @@ def test_9801(self, duckdb_cursor): assert x == (counter, words[counter % 4]) counter += 1 result = rel1.fetchmany(size=5) + + +ROW_COUNT = 1_000_000 +SMALL_BUFFER = "1MB" + + +def drain(res): + while res.fetchone() is not None: + pass + + +class TestStreamingSemantics: + """Which entry points stream, and what a stream promises once it is open.""" + + @pytest.fixture + def produced(self, duckdb_cursor): + counter = [0] + + def tally(i): + counter[0] += 1 + return i + + duckdb_cursor.execute(f"SET max_streaming_buffer_size='{SMALL_BUFFER}'") + duckdb_cursor.create_function("tally", tally, [BIGINT], BIGINT) + return counter + + TALLY_QUERY = f"SELECT tally(i) AS i FROM range({ROW_COUNT}) t(i)" + + def test_relation_row_fetch_streams(self, duckdb_cursor, produced): + res = duckdb_cursor.sql(self.TALLY_QUERY) + assert res.fetchone() == (0,) + res.close() + assert produced[0] < ROW_COUNT // 2 + + def test_relation_chunk_fetch_streams(self, duckdb_cursor, produced): + res = duckdb_cursor.sql(self.TALLY_QUERY) + assert len(res.fetch_df_chunk()) == duckdb.__standard_vector_size__ + res.close() + assert produced[0] < ROW_COUNT // 2 + + def test_connection_result_streams(self, duckdb_cursor, produced): + duckdb_cursor.execute(self.TALLY_QUERY) + assert duckdb_cursor.fetchone() == (0,) + duckdb_cursor.execute("SELECT 1") + assert produced[0] < ROW_COUNT // 2 + + def test_arrow_reader_streams(self, duckdb_cursor, produced): + pytest.importorskip("pyarrow") + reader = duckdb_cursor.sql(self.TALLY_QUERY).to_arrow_reader(1024) + assert len(reader.read_next_batch()) == 1024 + del reader + assert produced[0] < ROW_COUNT // 2 + + def test_arrow_capsule_streams(self, duckdb_cursor, produced): + pa = pytest.importorskip("pyarrow") + # The capsule has no batch size parameter and its batches hold a million rows + query = f"SELECT tally(i) AS i FROM range({2 * ROW_COUNT}) t(i)" + capsule = duckdb_cursor.sql(query).__arrow_c_stream__() + reader = pa.RecordBatchReader._import_from_c_capsule(capsule) + assert len(reader.read_next_batch()) > 0 + del reader + assert produced[0] < ROW_COUNT + ROW_COUNT // 2 + + def test_arrow_reader_after_row_fetch_returns_the_remainder(self, duckdb_cursor): + pytest.importorskip("pyarrow") + res = duckdb_cursor.sql("SELECT i FROM range(10) t(i)") + assert res.fetchone() == (0,) + assert res.to_arrow_reader().read_all().column("i").to_pylist() == list(range(1, 10)) + + def test_arrow_reader_over_a_retained_result(self, duckdb_cursor): + pytest.importorskip("pyarrow") + res = duckdb_cursor.sql("SELECT i FROM range(10) t(i)").execute() + assert res.to_arrow_reader(4).read_all().column("i").to_pylist() == list(range(10)) + + def test_second_statement_ends_the_open_arrow_reader(self, duckdb_cursor): + pytest.importorskip("pyarrow") + duckdb_cursor.execute(f"SET max_streaming_buffer_size='{SMALL_BUFFER}'") + reader = duckdb_cursor.sql(f"SELECT i FROM range({ROW_COUNT}) t(i)").to_arrow_reader(1024) + assert len(reader.read_next_batch()) == 1024 + + duckdb_cursor.execute("SELECT 42") + + # The engine's blocking fetch reports the ended query as closed rather than as cancelled + with pytest.raises(OSError, match=r"cancelled|closed query result"): + reader.read_all() + + def test_relation_whole_fetch_runs_to_the_end(self, duckdb_cursor, produced): + assert len(duckdb_cursor.sql(self.TALLY_QUERY).fetchall()) == ROW_COUNT + assert produced[0] == ROW_COUNT + + def test_whole_fetch_after_partial_fetch_returns_the_remainder(self, duckdb_cursor): + query = "SELECT i FROM range(10) t(i)" + + res = duckdb_cursor.sql(query) + assert res.fetchone() == (0,) + assert res.fetchall() == [(i,) for i in range(1, 10)] + + res = duckdb_cursor.sql(query) + assert res.fetchmany(3) == [(0,), (1,), (2,)] + assert res.df()["i"].tolist() == list(range(3, 10)) + + res = duckdb_cursor.sql(query) + assert res.fetchone() == (0,) + assert len(res.fetch_df_chunk(0)) == 0 + assert res.fetch_df_chunk()["i"].tolist() == list(range(1, 10)) + + res = duckdb_cursor.sql(query).execute() + assert res.fetchone() == (0,) + assert res.fetchall() == [(i,) for i in range(1, 10)] + + res = duckdb_cursor.sql(query).execute() + assert res.fetchone() == (0,) + assert res.df()["i"].tolist() == list(range(1, 10)) + + duckdb_cursor.execute(query) + assert duckdb_cursor.fetchone() == (0,) + assert duckdb_cursor.df()["i"].tolist() == list(range(1, 10)) + + def test_error_surfaces_on_the_fetch_that_reaches_it(self, duckdb_cursor): + duckdb_cursor.execute(f"SET max_streaming_buffer_size='{SMALL_BUFFER}'") + res = duckdb_cursor.sql( + f""" + SELECT CASE WHEN i < {ROW_COUNT} THEN i ELSE concat('hello', i::VARCHAR)::INT END AS i + FROM range({ROW_COUNT} + 1) t(i) + """ + ) + assert res.fetchone() == (0,) + with pytest.raises(duckdb.ConversionException): + drain(res) + + def test_second_statement_ends_the_open_stream(self, duckdb_cursor): + duckdb_cursor.execute(f"SET max_streaming_buffer_size='{SMALL_BUFFER}'") + res = duckdb_cursor.sql(f"SELECT i FROM range({ROW_COUNT}) t(i)") + assert res.fetchone() == (0,) + + duckdb_cursor.execute("SELECT 42") + + with pytest.raises(duckdb.InterruptException, match="cancelled"): + drain(res) + + def test_printing_ends_the_open_stream(self, duckdb_cursor): + duckdb_cursor.execute(f"SET max_streaming_buffer_size='{SMALL_BUFFER}'") + res = duckdb_cursor.sql(f"SELECT i FROM range({ROW_COUNT}) t(i)") + assert res.fetchone() == (0,) + + assert "i" in str(res) + + with pytest.raises(duckdb.InterruptException, match="cancelled"): + drain(res) + + def test_side_effecting_statement_falls_back_to_retained(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + + res = duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3) RETURNING i") + assert res.fetchone() == (1,) + assert res.fetchall() == [(2,), (3,)] + assert duckdb_cursor.execute("SELECT count(*) FROM t").fetchone() == (3,) + + def test_stream_outlives_the_connection(self): + con = duckdb.connect() + con.execute(f"SET max_streaming_buffer_size='{SMALL_BUFFER}'") + res = con.sql(f"SELECT i FROM range({ROW_COUNT}) t(i)") + assert res.fetchone() == (0,) + + del con + gc.collect() + + assert res.fetchall()[-1] == (ROW_COUNT - 1,) + + @pytest.mark.timeout(60) + def test_dropping_a_running_stream_does_not_deadlock(self, duckdb_cursor): + # Ending a query waits for its running tasks. A task inside a Python UDF cannot finish while + # the dropping thread holds the GIL, so dropping the result must release it first. + def identity(i): + return i + + duckdb_cursor.execute(f"SET max_streaming_buffer_size='{SMALL_BUFFER}'") + duckdb_cursor.create_function("identity", identity, [BIGINT], BIGINT) + res = duckdb_cursor.sql(f"SELECT identity(i) AS i FROM range({2 * ROW_COUNT}) t(i)") + assert res.fetchone() == (0,) + + del res + gc.collect() + + assert duckdb_cursor.execute("SELECT 42").fetchall() == [(42,)] diff --git a/tests/fast/arrow/test_arrow_fetch_recordbatch.py b/tests/fast/arrow/test_arrow_fetch_recordbatch.py index d060659f..2bb25874 100644 --- a/tests/fast/arrow/test_arrow_fetch_recordbatch.py +++ b/tests/fast/arrow/test_arrow_fetch_recordbatch.py @@ -250,9 +250,10 @@ def test_record_coverage(self, duckdb_cursor): def test_record_batch_query_error(self): duckdb_cursor = duckdb.connect() duckdb_cursor.execute("CREATE table t as select 'foo' as a;") + # A SELECT streams, so its execution error is reported by the first fetch + res = duckdb_cursor.execute("SELECT cast(a as double) FROM t") with pytest.raises(duckdb.ConversionException, match="Conversion Error"): - # 'execute' materializes the result, causing the error directly - duckdb_cursor.execute("SELECT cast(a as double) FROM t") + res.fetchall() def test_many_list_batches(self): conn = duckdb.connect() diff --git a/tests/fast/arrow/test_arrow_replacement_scan.py b/tests/fast/arrow/test_arrow_replacement_scan.py index 8c372a22..08d6cabe 100644 --- a/tests/fast/arrow/test_arrow_replacement_scan.py +++ b/tests/fast/arrow/test_arrow_replacement_scan.py @@ -1,3 +1,5 @@ +import gc +import weakref from pathlib import Path import pytest @@ -70,3 +72,102 @@ def test_arrow_dataset_replacement_scan(self, duckdb_cursor): con = duckdb.connect() assert con.execute("select count(*) from userdata_parquet_dataset").fetchone() == (1000,) + + +def _make_table(): + return pa.table({"a": [1, 2, 3, 4, 5]}) + + +def _make_dataset(): + return ds.dataset(_make_table()) + + +def _make_scanner(): + return ds.dataset(_make_table()).scanner() + + +def _make_record_batch_reader(): + table = _make_table() + return pa.RecordBatchReader.from_batches(table.schema, table.to_batches()) + + +def _make_polars_frame(): + pl = pytest.importorskip("polars") + return pl.DataFrame({"a": [1, 2, 3, 4, 5]}) + + +def _make_polars_lazy_frame(): + pl = pytest.importorskip("polars") + return pl.DataFrame({"a": [1, 2, 3, 4, 5]}).lazy() + + +FACTORIES = [ + ("table", _make_table), + ("dataset", _make_dataset), + ("scanner", _make_scanner), + ("record_batch_reader", _make_record_batch_reader), + ("polars_frame", _make_polars_frame), + ("polars_lazy_frame", _make_polars_lazy_frame), +] + +# An eager polars DataFrame is materialized through .to_arrow() before the scan is built, so the +# scan retains that conversion and not the frame itself. +RETAINED_FACTORIES = [entry for entry in FACTORIES if entry[0] != "polars_frame"] + + +class TestArrowScanFactoryOwnership: + """The scan factory owns the Arrow object for as long as the engine holds the scan.""" + + @pytest.mark.parametrize(("label", "make"), FACTORIES, ids=[label for label, _ in FACTORIES]) + def test_scanned_object_survives_its_python_variable(self, label, make): + con = duckdb.connect() + scanned = make() + rel = con.sql("select sum(a) as s from scanned") + + del scanned + gc.collect() + + assert rel.fetchall() == [(15,)] + + @pytest.mark.parametrize(("label", "make"), RETAINED_FACTORIES, ids=[label for label, _ in RETAINED_FACTORIES]) + def test_scanned_object_is_released_with_the_relation(self, label, make): + con = duckdb.connect() + + # The scan reads the caller's frame locals, and before Python 3.13 that read caches a dict on + # the frame which keeps the object alive until the frame returns, so the scan gets its own frame. + def scan(con): + scanned = make() + rel = con.sql("select sum(a) as s from scanned") + return weakref.ref(scanned), rel + + ref, rel = scan(con) + gc.collect() + assert ref() is not None + + assert rel.fetchall() == [(15,)] + + del rel + con.close() + del con + gc.collect() + assert ref() is None + + @pytest.mark.parametrize(("label", "make"), FACTORIES, ids=[label for label, _ in FACTORIES]) + def test_scan_through_execute_outlives_the_bind(self, label, make): + # Without a relation there is no dependency to keep the scan alive past binding: the bind + # data alone must hold the factory until the scan runs. + con = duckdb.connect() + scanned = make() + assert con.execute("select sum(a) as s from scanned").fetchall() == [(15,)] + del scanned + + def test_repeated_scans_of_one_object(self): + con = duckdb.connect() + scanned = _make_table() + rel = con.sql("select sum(a) as s from scanned") + + del scanned + gc.collect() + + assert rel.fetchall() == [(15,)] + assert rel.fetchall() == [(15,)] diff --git a/tests/fast/pandas/test_implicit_pandas_scan.py b/tests/fast/pandas/test_implicit_pandas_scan.py index 65fd3da9..d8040ab2 100644 --- a/tests/fast/pandas/test_implicit_pandas_scan.py +++ b/tests/fast/pandas/test_implicit_pandas_scan.py @@ -1,10 +1,25 @@ # simple DB API testcase +import pytest + import duckdb import pandas as pd class TestImplicitPandasScan: + def test_pandas_scan_takes_no_sql_arguments(self, duckdb_cursor): + # The dataframe travels through the bind input, never through a value SQL text can forge + with pytest.raises(duckdb.BinderException, match="No function matches"): + duckdb_cursor.execute("select * from pandas_scan(140234567890)") + with pytest.raises(duckdb.BinderException, match="requires a dataframe bind input"): + duckdb_cursor.execute("select * from pandas_scan()") + + def test_two_dataframes_in_one_query(self, duckdb_cursor): + lhs = pd.DataFrame({"k": [1, 2, 3], "a": ["x", "y", "z"]}) # noqa: F841 + rhs = pd.DataFrame({"k": [2, 3, 4], "b": [20, 30, 40]}) # noqa: F841 + rows = duckdb_cursor.execute("select a, b from lhs join rhs using (k) order by k").fetchall() + assert rows == [("y", 20), ("z", 30)] + def test_local_pandas_scan(self, duckdb_cursor): con = duckdb.connect() df = pd.DataFrame([{"COL1": "val1", "CoL2": 1.05}, {"COL1": "val3", "CoL2": 17}]) # noqa: F841 diff --git a/tests/fast/pandas/test_pandas_category.py b/tests/fast/pandas/test_pandas_category.py index b5815cf6..a44c671a 100644 --- a/tests/fast/pandas/test_pandas_category.py +++ b/tests/fast/pandas/test_pandas_category.py @@ -109,7 +109,7 @@ def test_empty_categorical(self, duckdb_cursor): res = duckdb_cursor.table("test").fetchall() assert res == [] - with pytest.raises(duckdb.ConversionException, match="Could not convert string 'test' to UINT8"): + with pytest.raises(duckdb.ConversionException, match="Could not convert string 'test' to ENUM"): duckdb_cursor.execute("insert into test VALUES('test')") duckdb_cursor.execute("insert into test VALUES(NULL)") res = duckdb_cursor.table("test").fetchall() diff --git a/tests/fast/pandas/test_pandas_enum.py b/tests/fast/pandas/test_pandas_enum.py index 298adb8f..5070fa1d 100644 --- a/tests/fast/pandas/test_pandas_enum.py +++ b/tests/fast/pandas/test_pandas_enum.py @@ -40,8 +40,7 @@ def test_3479(self, duckdb_cursor): duckdb_cursor.register("df", df) with pytest.raises( duckdb.ConversionException, - match="Type UINT8 with value 0 can't be cast because the value is out of range for the destination " - "type UINT8", + match="with value berlioz can't be cast to the destination type ENUM", ): duckdb_cursor.execute("INSERT INTO tab SELECT * FROM df;") diff --git a/tests/fast/test_map.py b/tests/fast/test_map.py index f3d6f16d..a0ebeb18 100644 --- a/tests/fast/test_map.py +++ b/tests/fast/test_map.py @@ -1,3 +1,4 @@ +import gc import re from datetime import date, timedelta from typing import NoReturn @@ -49,6 +50,23 @@ def capture(df): alive = sum(1 for r in refs if r() is not None) assert alive == 0, f"{alive}/{len(refs)} per-chunk input DataFrames leaked (pinned by arg tuple)" + def test_map_function_takes_no_sql_arguments(self, duckdb_cursor): + # The callable travels through the bind input, never through a value SQL text can forge + with pytest.raises(duckdb.BinderException, match="No function matches"): + duckdb_cursor.sql("select * from python_map_function(1, 2, 3)") + with pytest.raises(duckdb.BinderException, match="requires a callable bind input"): + duckdb_cursor.sql("select * from python_map_function((select 1))") + + def test_mapped_view_outlives_the_relation(self, duckdb_cursor): + def double(df): + return df.assign(i=df["i"] * 2) + + rel = duckdb_cursor.sql("select 21 as i").map(double) + rel.create_view("mapped") + del rel + gc.collect() + assert duckdb_cursor.sql("select i from mapped").fetchall() == [(42,)] + def test_evil_map(self, duckdb_cursor): testrel = duckdb.values([1, 2]) rel = testrel.map(evil1, schema={"i": str}) diff --git a/tests/fast/test_relation.py b/tests/fast/test_relation.py index 7c47c391..666db4cd 100644 --- a/tests/fast/test_relation.py +++ b/tests/fast/test_relation.py @@ -677,6 +677,34 @@ def create_view(con, view_name: str) -> None: res = con.sql("select * from vw").fetchall() assert res == expected + @pytest.mark.parametrize( + "make_relation", + [ + lambda con: con.from_df(pd.DataFrame({"a": [1, 2, 3]})), + lambda con: con.from_arrow(pytest.importorskip("pyarrow").table({"a": [1, 2, 3]})), + lambda con: con.sql("select range as a from range(3)").map(lambda df: df), + ], + ids=["pandas", "arrow", "map"], + ) + def test_persistent_view_over_python_object(self, tmp_path, make_relation): + # A persistent view is written to the database file, and a scan over a Python object has no + # representation there, so the catalog change must be refused rather than stored dangling. + con = duckdb.connect(str(tmp_path / "db.duckdb")) + rel = make_relation(con) + with pytest.raises(duckdb.TransactionException, match="Cannot serialize a table function"): + rel.create_view("vw") + assert con.sql("select view_name from duckdb_views() where not internal").fetchall() == [] + + con.execute("begin") + rel.create_view("vw") + with pytest.raises(duckdb.TransactionException, match="Cannot serialize a table function"): + con.execute("commit") + assert con.sql("select view_name from duckdb_views() where not internal").fetchall() == [] + + # Temporary views never reach the file, so they keep working on the same database + assert rel.query("tmp_vw", "select count(*) from tmp_vw").fetchall() == [(3,)] + con.close() + def test_relation_select_dtypes_quotes_identifiers_with_spaces(self, duckdb_cursor): df = pd.DataFrame({"na me": ["alice", "bob"], "x": [1, 2]}) rel = duckdb_cursor.from_df(df) diff --git a/tests/fast/test_variant.py b/tests/fast/test_variant.py index af4496a9..33daada7 100644 --- a/tests/fast/test_variant.py +++ b/tests/fast/test_variant.py @@ -143,16 +143,17 @@ def test_null_handling(self): class TestVariantArrow: """Tests for Arrow/Polars — blocked on DuckDB core Arrow support.""" - @pytest.mark.xfail(strict=True, reason="Arrow export for VARIANT not yet supported in DuckDB core") def test_to_arrow_table(self): + pytest.importorskip("pyarrow") duckdb.sql("SELECT 42::VARIANT AS v").arrow() @pytest.mark.xfail(strict=True, reason="Arrow export for VARIANT not yet supported in DuckDB core") def test_fetch_arrow_reader(self): duckdb.sql("SELECT 42::VARIANT AS v").fetch_arrow_reader() - @pytest.mark.xfail(strict=True, reason="Polars uses Arrow, which doesn't support VARIANT yet") def test_polars(self): + pytest.importorskip("pyarrow") + pytest.importorskip("polars") duckdb.sql("SELECT 42::VARIANT AS v").pl() diff --git a/tests/fast/udf/test_scalar.py b/tests/fast/udf/test_scalar.py index 348b8eb1..9465fbe0 100644 --- a/tests/fast/udf/test_scalar.py +++ b/tests/fast/udf/test_scalar.py @@ -311,7 +311,9 @@ def func(x: int) -> int: con.create_function("func", func) - rel.fetchall() + # Registering the function ended the open query, and the connection is usable afterwards + with pytest.raises(duckdb.InterruptException, match="cancelled"): + rel.fetchall() res = con.sql("select func(5)").fetchall() assert res == [(5,)] diff --git a/tests/fast/udf/test_transactionality.py b/tests/fast/udf/test_transactionality.py index acad21ef..b5428146 100644 --- a/tests/fast/udf/test_transactionality.py +++ b/tests/fast/udf/test_transactionality.py @@ -4,7 +4,6 @@ class TestUDFTransactionality: - @pytest.mark.xfail(reason="fetchone() does not realize the stream result was closed before completion") def test_type_coverage(self, duckdb_cursor): rel = duckdb_cursor.sql("select * from range(4096)") res = rel.fetchone() @@ -15,5 +14,5 @@ def my_func(x: str) -> int: duckdb_cursor.create_function("test", my_func) - with pytest.raises(duckdb.InvalidInputException, match="result closed"): - res = rel.fetchone() + with pytest.raises(duckdb.InterruptException, match="cancelled"): + rel.fetchone()