From 7fa2cfe73c8b904cfcf8da0e3a16fdfeb53fb254 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Tue, 10 Nov 2026 17:50:30 +0000 Subject: [PATCH 1/6] Hold one connection claim from submission to settling in execute, executemany and the DB-API Arrow fetch --- src/engine/connection.cpp | 29 +-- src/engine/result.cpp | 26 ++- .../duckdb_python/engine/connection.hpp | 18 +- src/include/duckdb_python/engine/result.hpp | 22 ++- .../pyconnection/pyconnection.hpp | 92 ++++++--- src/include/duckdb_python/pyresult.hpp | 15 +- src/pyconnection.cpp | 178 ++++++++++-------- src/pyresult.cpp | 74 +++++--- tests/fast/api/test_appender.py | 77 ++++++++ tests/fast/api/test_execute_contract.py | 158 +++++++++++++++- 10 files changed, 516 insertions(+), 173 deletions(-) diff --git a/src/engine/connection.cpp b/src/engine/connection.cpp index d9733108..1f98f194 100644 --- a/src/engine/connection.cpp +++ b/src/engine/connection.cpp @@ -65,12 +65,16 @@ bool CallbackScope::Inside(const ClientContext &context) { return false; } +void ThrowBusy() { + throw ResourceInUseException(BUSY_ERROR); +} + ConnectionClaim::ConnectionClaim(ConnectionState &state_p, const ClientContext &context) : state(state_p) { if (CallbackScope::Inside(context)) { throw ResourceInUseException(CALLBACK_ERROR); } if (state.busy.exchange(true)) { - throw ResourceInUseException(BUSY_ERROR); + ThrowBusy(); } if (Exiting()) { state.admitted_during_exit = OnExitThread(); @@ -89,7 +93,7 @@ void ConnectionState::ThrowIfBusy(const ClientContext &context) const { throw ResourceInUseException(CALLBACK_ERROR); } if (busy) { - throw ResourceInUseException(BUSY_ERROR); + ThrowBusy(); } } @@ -340,12 +344,6 @@ ConnectionClaim Session::Claim() { return ConnectionClaim(*state, *context); } -shared_ptr Session::Submit(unique_ptr statement, identifier_map_t &values, - const Format &format) { - auto claim = Claim(); - return SubmitClaimed(std::move(statement), values, format); -} - shared_ptr Session::SubmitClaimed(unique_ptr statement, identifier_map_t &values, const Format &format) { auto written_type = statement->type; @@ -357,9 +355,8 @@ shared_ptr Session::SubmitClaimed(unique_ptr statement, return result; } -shared_ptr Session::Submit(Prepared &prepared, identifier_map_t &values, - const Format &format) { - ConnectionClaim claim(*state, *context); +shared_ptr Session::SubmitClaimed(Prepared &prepared, identifier_map_t &values, + const Format &format) { registry->CountStatement(prepared.statement->GetStatementType()); unique_ptr handle; { @@ -372,8 +369,8 @@ shared_ptr Session::Submit(Prepared &prepared, identifier_map_t Session::Submit(unique_ptr collection, const vector &names, - const Format &format) { +shared_ptr Session::SubmitClaimed(unique_ptr collection, const vector &names, + const Format &format) { // The binder rejects duplicate column names; the caller keeps the originals auto deduplicated_names = names; QueryResult::DeduplicateColumns(deduplicated_names); @@ -386,7 +383,7 @@ shared_ptr Session::Submit(unique_ptr collection, auto select = make_uniq(); select->node = std::move(select_node); identifier_map_t no_values; - return Submit(std::move(select), no_values, format); + return SubmitClaimed(std::move(select), no_values, format); } StatementSignature Session::Bind(unique_ptr statement, ResolvedTables &tables) { @@ -557,6 +554,10 @@ shared_ptr StatementGroup::Choose(shared_ptr last) { void StatementGroup::RefuseParameters() { auto claim = session.Claim(); + RefuseParametersClaimed(); +} + +void StatementGroup::RefuseParametersClaimed() { Fail(ErrorData(InvalidInputException("Parameters were given, but the last statement of the query takes none"))); } diff --git a/src/engine/result.cpp b/src/engine/result.cpp index d2350284..aae5de4a 100644 --- a/src/engine/result.cpp +++ b/src/engine/result.cpp @@ -218,12 +218,6 @@ shared_ptr Result::TryGetContext() const { return context; } -void Result::Retain(const InterruptCheck &check) { - DUCKDB_PY_ASSERT_GIL_RELEASED(); - auto claim = ClaimWhileRunning(); - RetainClaimed(check); -} - void Result::CloseAfterFailure(std::exception_ptr error) noexcept { DUCKDB_PY_ASSERT_GIL_RELEASED(); { @@ -281,12 +275,6 @@ void Result::RetainClaimed(const InterruptCheck &check) { state = State::RETAINED; } -void Result::Settle(const InterruptCheck &check) { - DUCKDB_PY_ASSERT_GIL_RELEASED(); - auto claim = ClaimWhileRunning(); - SettleClaimed(check); -} - void Result::SettleClaimed(const InterruptCheck &check) { DUCKDB_PY_ASSERT_GIL_RELEASED(); if (state != State::PENDING) { @@ -375,10 +363,16 @@ unique_ptr Result::FetchArray(const InterruptCheck &check) { unique_ptr Result::TakeChunks(const InterruptCheck &check) { DUCKDB_PY_ASSERT_GIL_RELEASED(); - D_ASSERT(!format.IsArrow()); ThrowIfFailed(); // Claimed only while the query can still run engine work: rows already in the client need no connection auto claim = ClaimWhileRunning(); + return TakeChunksClaimed(check); +} + +unique_ptr Result::TakeChunksClaimed(const InterruptCheck &check) { + DUCKDB_PY_ASSERT_GIL_RELEASED(); + D_ASSERT(!format.IsArrow()); + ThrowIfFailed(); if (state == State::PENDING) { RetainClaimed(check); } @@ -448,9 +442,13 @@ void Result::CopyArrowSchema(ArrowSchema &out) const { } void Result::BuildArrowSchema(const vector &schema_names, ArrowSchema &out) const { + ConnectionClaim claim(*connection_state, *TryGetContext()); + BuildArrowSchemaClaimed(schema_names, out); +} + +void Result::BuildArrowSchemaClaimed(const vector &schema_names, ArrowSchema &out) const { auto context = TryGetContext(); auto properties = client_properties; - ConnectionClaim claim(*connection_state, *context); context->RunFunctionInTransaction( [&]() { ArrowConverter::ToArrowSchema(&out, Types(), schema_names, properties); }); } diff --git a/src/include/duckdb_python/engine/connection.hpp b/src/include/duckdb_python/engine/connection.hpp index f1f35be0..be4c9687 100644 --- a/src/include/duckdb_python/engine/connection.hpp +++ b/src/include/duckdb_python/engine/connection.hpp @@ -168,6 +168,9 @@ class SourceRegistry : public ObjectCacheEntry { std::deque released_order; }; +//! The refusal of a call that finds the connection taken by another call +[[noreturn]] void ThrowBusy(); + //! Holds the connection for one stretch of engine work. A second claim, and any claim from a callback of the //! connection's queries, is refused instead of waiting: the work it would wait for may be waiting for it. class ConnectionClaim { @@ -205,15 +208,16 @@ class Session { return *context; } Prepared Prepare(unique_ptr statement); - shared_ptr Submit(unique_ptr statement, identifier_map_t &values, - const Format &format); - //! Submit for a caller that already holds the connection's claim + //! Submits a statement whose result the caller reads later, for a caller that holds the connection's claim; + //! the caller keeps its claim until the result is settled or retained, so no other call finds the statement + //! pending. Run and RunClaimed below submit and complete a statement in one go. shared_ptr SubmitClaimed(unique_ptr statement, identifier_map_t &values, const Format &format); - shared_ptr Submit(Prepared &prepared, identifier_map_t &values, const Format &format); + shared_ptr SubmitClaimed(Prepared &prepared, identifier_map_t &values, + const Format &format); //! Scans a collection, for example to read rows a result already holds in another format - shared_ptr Submit(unique_ptr collection, const vector &names, - const Format &format); + shared_ptr SubmitClaimed(unique_ptr collection, const vector &names, + const Format &format); //! Binds a statement without running it, keeping the tables its replacement scans resolved in tables StatementSignature Bind(unique_ptr statement, ResolvedTables &tables); //! Binds a statement with values for its parameters, as the engine rebinds a prepared statement for the values it @@ -267,6 +271,8 @@ class StatementGroup { shared_ptr Choose(shared_ptr last); //! Throws for parameters given to a group whose text after the statements that ran held no statement [[noreturn]] void RefuseParameters(); + //! RefuseParameters for a caller that already holds the connection's claim + [[noreturn]] void RefuseParametersClaimed(); private: bool HasMore(); diff --git a/src/include/duckdb_python/engine/result.hpp b/src/include/duckdb_python/engine/result.hpp index feb5cadd..ca20b063 100644 --- a/src/include/duckdb_python/engine/result.hpp +++ b/src/include/duckdb_python/engine/result.hpp @@ -52,9 +52,10 @@ class Format { class ConnectionState; class ConnectionClaim; -//! A submitted statement. Retention is settled by the first consumer: a fetch opens a stream, Retain and the -//! Take calls keep every row. Every call that drives the query claims the connection, takes the client context -//! lock and may wait for running tasks, so callers must not hold a lock those tasks need. +//! A submitted statement. Retention is settled by the first consumer: a fetch opens a stream, RetainClaimed and +//! the Take calls keep every row. Every call that drives the query runs under the connection's claim, its own or +//! the caller's, takes the client context lock and may wait for running tasks, so callers must not hold a lock +//! those tasks need. class Result : public enable_shared_from_this { public: //! written_type is the submitted statement's own type, since the handle of a prepared execution reports EXECUTE @@ -105,14 +106,11 @@ class Result : public enable_shared_from_this { shared_ptr TryGetContext() const; //! The database whose buffers hold the rows, null once it is gone. Weak for the same reason as the context. shared_ptr TryGetDatabase() const; - //! Keeps every row: drives the query to its end unless it already ended. A failure ends the query. - void Retain(const InterruptCheck &check); - //! Retain for a caller that already holds the connection's claim. A failure ends the query. + //! Keeps every row: drives the query to its end unless it already ended, for a caller that holds the + //! connection's claim. A failure ends the query. void RetainClaimed(const InterruptCheck &check); - //! Runs the statement until its first unit is ready or it ends, without deciding how its rows are read, - //! and throws the error it ended with. A failure ends the query. - void Settle(const InterruptCheck &check); - //! Settle for a caller that already holds the connection's claim. A failure ends the query. + //! Runs the statement until its first unit is ready or it ends, without deciding how its rows are read, and + //! throws the error it ended with, for a caller that holds the connection's claim. A failure ends the query. void SettleClaimed(const InterruptCheck &check); //! Rows of a retained result. Materializes it first. idx_t RowCount(const InterruptCheck &check); @@ -122,12 +120,16 @@ class Result : public enable_shared_from_this { unique_ptr FetchArray(const InterruptCheck &check); //! Every remaining row, the result is empty afterwards unique_ptr TakeChunks(const InterruptCheck &check); + //! TakeChunks for a caller that already holds the connection's claim + unique_ptr TakeChunksClaimed(const InterruptCheck &check); //! Every remaining array, the result is empty afterwards vector> TakeArrays(const InterruptCheck &check); //! A copy of the result's Arrow schema, owned by the caller. Arrow results only. void CopyArrowSchema(ArrowSchema &out) const; //! An Arrow schema for the result's types under the given names, built in a transaction on its connection void BuildArrowSchema(const vector &names, ArrowSchema &out) const; + //! BuildArrowSchema for a caller that already holds the connection's claim + void BuildArrowSchemaClaimed(const vector &names, ArrowSchema &out) const; //! Renders the rows of a retained result string ToBox(BoxRendererContext &context, const BoxRendererConfig &config); diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 77953702..03a85cf4 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -88,8 +88,9 @@ class DbapiResult { void Supersede(); //! Every row the result still holds, under names, which leaves it read out as a consuming fetch does; null //! when the rows were used up already. The rows are gone even when the take fails, since it ran the - //! statement to its end or ended it. - unique_ptr TakeRows(vector &names); + //! statement to its end or ended it. The caller holds the connection's claim, so nothing can refuse the take + //! once it has begun. + unique_ptr TakeRowsClaimed(vector &names); shared_ptr EngineResult() const { return result ? result->EngineResult() : nullptr; } @@ -266,38 +267,90 @@ struct DuckDBPyConnection : public std::enable_shared_from_this lock_; + bool marks_submitting_call = false; }; ConnectionGuard con; Cursors cursors; - // Recursive so that the outer lock taken at the top of execute/fetch - // methods (while still holding the GIL) does not deadlock against the - // inner lock taken by PrepareQuery / ExecuteInternal / - // PrepareAndSubmitInternal (after releasing the GIL). Serialises every - // path that touches `con.result` so concurrent calls on a single - // DuckDBPyConnection cannot dereference an already-freed result — see - // duckdb-python#435. + //! Serialises every path that touches con.result, so concurrent calls on one connection cannot dereference a + //! freed result (duckdb-python#435). Recursive, so the lock taken at the top of execute and the fetches does not + //! deadlock against the one PrepareQuery and ExecuteInternal take after releasing the GIL std::recursive_mutex py_connection_lock; + //! Whether the thread holding py_connection_lock is inside a call that submits a statement (IfHeld::REFUSE) + std::atomic mutex_held_by_submitting_call {false}; + //! How many times a submitting call has taken py_connection_lock + std::atomic submitting_entries {0}; //! The database's store of file-like objects for reading, kept once this connection has found it shared_ptr object_store; case_insensitive_map_t> registered_functions; @@ -372,9 +425,6 @@ struct DuckDBPyConnection : public std::enable_shared_from_this statement); shared_ptr ExecuteInternal(engine::Prepared &prep, nb::object params, engine::StatementGroup &statements); - //! Binds the parameters and submits the statement. The result is returned undriven. - shared_ptr PrepareAndSubmitInternal(unique_ptr statement, - nb::object params = nb::list()); engine::Session EngineSession(); std::shared_ptr Execute(const nb::object &query, nb::object params = nb::list()); diff --git a/src/include/duckdb_python/pyresult.hpp b/src/include/duckdb_python/pyresult.hpp index abad6e35..28188411 100644 --- a/src/include/duckdb_python/pyresult.hpp +++ b/src/include/duckdb_python/pyresult.hpp @@ -76,6 +76,12 @@ struct DuckDBPyResult { duckdb::pyarrow::Table FetchArrowTable(idx_t rows_per_batch, bool to_polars); duckdb::pyarrow::RecordBatchReader FetchRecordBatchReader(idx_t rows_per_batch = DEFAULT_ARROW_BATCH_SIZE); + //! The two fetches above over a schema built with BuildArrowSchemaClaimed, which they take over + duckdb::pyarrow::Table FetchArrowTable(ArrowSchema &schema); + duckdb::pyarrow::RecordBatchReader FetchRecordBatchReader(ArrowSchema &schema, idx_t rows_per_batch); + //! The result's Arrow schema under the names the Python layer reports, for a caller that holds the connection's + //! claim; out is the caller's to release until a fetch takes it over + void BuildArrowSchemaClaimed(ArrowSchema &out); nb::object FetchArrowCapsule(idx_t rows_per_batch = DEFAULT_ARROW_BATCH_SIZE); static void CheckBatchSize(idx_t rows_per_batch); @@ -89,8 +95,8 @@ struct DuckDBPyResult { unique_ptr FetchChunk(); //! Every row the result still holds, the rows buffered from a partial fetch first; the result is empty - //! afterwards. Call with the GIL held. - unique_ptr TakeRows(const engine::InterruptCheck &check); + //! afterwards. Call with the GIL held and the connection's claim. + unique_ptr TakeRowsClaimed(const engine::InterruptCheck &check); //! The names the Python layer reports, which are names_override when that is set. const vector &ResultNames() const; //! The rows a completed DML statement changed, -1 otherwise @@ -123,8 +129,11 @@ struct DuckDBPyResult { template T RunWithArrowSchema(const std::function &fun, bool dedup_col_names); + duckdb::pyarrow::Table ConvertToArrowTable(const ArrowSchema &schema); + duckdb::pyarrow::RecordBatchReader ImportRecordBatchReader(ArrowSchema *prebuilt, idx_t rows_per_batch); //! The stream's private data owns the result, and its callbacks run without the GIL. - ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch); + //! prebuilt, when given, becomes the stream's schema; otherwise the result's own is copied or built + ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch, ArrowSchema *prebuilt); private: shared_ptr result; diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index c82b65d1..12b1ce65 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -181,7 +181,7 @@ void DbapiResult::Supersede() { } } -unique_ptr DbapiResult::TakeRows(vector &names) { +unique_ptr DbapiResult::TakeRowsClaimed(vector &names) { if (!result) { return nullptr; } @@ -189,7 +189,7 @@ unique_ptr DbapiResult::TakeRows(vector &names // failure must leave the result read out rather than reporting its error at every later fetch auto taken = std::move(result); names = taken->ResultNames(); - return taken->TakeRows(DuckDBPyConnection::CheckSignals); + return taken->TakeRowsClaimed(DuckDBPyConnection::CheckSignals); } void DbapiResult::Abort() { @@ -627,7 +627,7 @@ void DuckDBPyConnection::SupersedeResult(bool cancel) { std::shared_ptr DuckDBPyConnection::ExecuteMany(const nb::object &query, nb::object params_p) { nb::gil_scoped_acquire gil; - ConnectionLockGuard conn_lock(*this); + ConnectionLockGuard conn_lock(*this, IfHeld::REFUSE); if (params_p.is_none()) { params_p = nb::list(); } @@ -768,6 +768,31 @@ identifier_map_t TransformPreparedParameters(ClientContext & return named_values; } +//! Submits through submit under the claim the caller holds, then runs the result to its first unit or, when +//! retain is set, to its end; a statement that expanded to nothing refuses the parameters it was given. A failure +//! ends the query here, with the GIL released, since ending a pending query waits for its tasks +template +static shared_ptr SubmitAndDrive(engine::StatementGroup &statements, bool parameters_given, bool retain, + SUBMIT &&submit) { + DUCKDB_PY_ASSERT_GIL_RELEASED(); + shared_ptr res = submit(); + try { + if (parameters_given && ExpandedToNothing(*res)) { + statements.RefuseParametersClaimed(); + } + statements.ReleaseIfReplaced(*res); + if (retain) { + res->RetainClaimed(DuckDBPyConnection::CheckSignals); + } else { + res->SettleClaimed(DuckDBPyConnection::CheckSignals); + } + } catch (...) { + res->Close(); + throw; + } + return res; +} + engine::Prepared DuckDBPyConnection::PrepareQuery(unique_ptr statement) { auto session = EngineSession(); engine::Prepared prep; @@ -801,38 +826,10 @@ shared_ptr DuckDBPyConnection::ExecuteInternal(engine::Prepared DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; unique_lock lock(py_connection_lock); - - res = session.Submit(prep, named_values, engine::Format::Chunks()); - if (parameters_given && ExpandedToNothing(*res)) { - try { - statements.RefuseParameters(); - } catch (...) { - // Ended here, with the GIL released, since ending the pending query waits for its tasks - res->Close(); - throw; - } - } - statements.ReleaseIfReplaced(*res); - res->Retain(CheckSignals); - } - return res; -} - -shared_ptr DuckDBPyConnection::PrepareAndSubmitInternal(unique_ptr statement, - nb::object params) { - if (params.is_none()) { - params = nb::list(); - } - auto session = EngineSession(); - - auto named_values = TransformPreparedParameters(session.Context(), params); - - shared_ptr res; - { - DUCKDB_PY_ASSERT_GIL_HELD(); - nb::gil_scoped_release release; - unique_lock lock(py_connection_lock); - res = session.Submit(std::move(statement), named_values, engine::Format::Chunks()); + // One claim from submission to the end of the run, so no other call finds the statement pending + auto claim = session.Claim(); + res = SubmitAndDrive(statements, parameters_given, true, + [&]() { return session.SubmitClaimed(prep, named_values, engine::Format::Chunks()); }); } return res; } @@ -875,7 +872,7 @@ std::shared_ptr DuckDBPyConnection::ExecuteFromString(const std::shared_ptr DuckDBPyConnection::Execute(const nb::object &query, nb::object params) { nb::gil_scoped_acquire gil; - ConnectionLockGuard conn_lock(*this); + ConnectionLockGuard conn_lock(*this, IfHeld::REFUSE); auto statements = GetStatements(query); auto parameters_given = HasParameters(params); @@ -886,27 +883,25 @@ std::shared_ptr DuckDBPyConnection::Execute(const nb::object } shared_ptr res; if (last_statement) { - res = PrepareAndSubmitInternal(std::move(last_statement), std::move(params)); - if (parameters_given && ExpandedToNothing(*res)) { - try { - RefuseParameters(statements); - } catch (...) { - // Ending the pending query waits for its tasks, so the GIL is released for it - duckdb::PyUtil::GilReleaseIfHeld release; - res->Close(); - throw; - } + if (params.is_none()) { + params = nb::list(); } - } - { + auto session = EngineSession(); + auto named_values = TransformPreparedParameters(session.Context(), params); DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; unique_lock lock(py_connection_lock); - if (res) { - statements.ReleaseIfReplaced(*res); - res->Settle(CheckSignals); - } + // One claim from submission to settling, so no other call finds the statement pending + auto claim = session.Claim(); + res = SubmitAndDrive(statements, parameters_given, false, [&]() { + return session.SubmitClaimed(std::move(last_statement), named_values, engine::Format::Chunks()); + }); res = statements.Choose(std::move(res)); + } else { + DUCKDB_PY_ASSERT_GIL_HELD(); + nb::gil_scoped_release release; + unique_lock lock(py_connection_lock); + res = statements.Choose(nullptr); } if (!res) { // TODO: should we throw? @@ -1513,8 +1508,19 @@ nb::dict DuckDBPyConnection::FetchTF() { return result.FetchTF(); } +//! Releases an Arrow schema that was never handed over, after a throw before the import into pyarrow, which takes the +//! schema over even when it fails +struct ArrowSchemaHolder { + ArrowSchema schema {}; + ~ArrowSchemaHolder() { + if (schema.release) { + schema.release(&schema); + } + } +}; + nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { - ConnectionLockGuard conn_lock(*this); + ConnectionLockGuard conn_lock(*this, IfHeld::REFUSE); // Asked before the result, so a closed connection reports itself and not the slot its close emptied auto session = EngineSession(); if (!con.HasResult()) { @@ -1534,41 +1540,53 @@ nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { // Imported before the engine checks, so no Python code runs between what they find and the take nb::module_::import_("pyarrow"); auto held = con.GetResult().EngineResult(); + unique_ptr converted; + ArrowSchemaHolder schema; { - DUCKDB_PY_ASSERT_GIL_HELD(); - nb::gil_scoped_release release; + // One claim from the checks through the take of the rows, the statement over them, its settle or retain and + // the schema build, so a refusal comes before the rows are taken and never loses them. The claim is a flag + // that never waits, so holding it while the GIL is held is safe auto claim = session.Claim(); - // The engine refuses an aborted transaction here under its own lock, and a running result keeps its - // transaction active, so the empty function neither begins nor commits anything the result depends on - session.Context().RunFunctionInTransaction([] {}); - // Another open result, a relation's reader for one, would refuse the statement over the rows after the take; - // a running DB-API result is itself the one open result and the take drains it - if (held && !held->Running()) { - engine::VerifyNoOpenResult(session.Context()); + { + DUCKDB_PY_ASSERT_GIL_HELD(); + nb::gil_scoped_release release; + // The engine refuses an aborted transaction here under its own lock, and a running result keeps its + // transaction active, so the empty function neither begins nor commits anything the result depends on + session.Context().RunFunctionInTransaction([] {}); + // Another open result, a relation's reader for one, would refuse the statement over the rows after the + // take; a running DB-API result is itself the one open result and the take drains it + if (held && !held->Running()) { + engine::VerifyNoOpenResult(session.Context()); + } } - } - vector names; - auto rows = con.GetResult().TakeRows(names); - if (!rows) { - return nb::none(); - } - shared_ptr arrow; - { - DUCKDB_PY_ASSERT_GIL_HELD(); - nb::gil_scoped_release release; - arrow = session.Submit(std::move(rows), names, format); - if (reader) { - arrow->Settle(DuckDBPyConnection::CheckSignals); - } else { - arrow->Retain(DuckDBPyConnection::CheckSignals); + vector names; + auto rows = con.GetResult().TakeRowsClaimed(names); + if (!rows) { + return nb::none(); + } + shared_ptr arrow; + { + DUCKDB_PY_ASSERT_GIL_HELD(); + nb::gil_scoped_release release; + arrow = session.SubmitClaimed(std::move(rows), names, format); + if (reader) { + arrow->SettleClaimed(DuckDBPyConnection::CheckSignals); + } else { + arrow->RetainClaimed(DuckDBPyConnection::CheckSignals); + } } + converted = make_uniq(std::move(arrow), std::move(names)); + // The converter reads the connection's settings and its catalog's types, so the schema is built here + converted->BuildArrowSchemaClaimed(schema.schema); } - DuckDBPyResult converted(std::move(arrow), std::move(names)); + // The rows are retained, or the stream is settled, so the conversion into pyarrow holds the connection no longer + // and refuses no other call + conn_lock.DropSubmittingMark(); if (reader) { - return nb::cast(converted.FetchRecordBatchReader(batch_size)); + return nb::cast(converted->FetchRecordBatchReader(schema.schema, batch_size)); } - return nb::cast(converted.FetchArrowTable(batch_size, false)); + return nb::cast(converted->FetchArrowTable(schema.schema)); } identifier_map_t TransformPyConfigDict(const nb::dict &py_config_dict) { diff --git a/src/pyresult.cpp b/src/pyresult.cpp index b796af6b..74b8c5ab 100644 --- a/src/pyresult.cpp +++ b/src/pyresult.cpp @@ -114,16 +114,16 @@ unique_ptr DuckDBPyResult::TakeBufferedRows() { return remainder; } -unique_ptr DuckDBPyResult::TakeRows(const engine::InterruptCheck &check) { +unique_ptr DuckDBPyResult::TakeRowsClaimed(const engine::InterruptCheck &check) { if (Empty()) { - throw InternalException("TakeRows called without a result object"); + throw InternalException("TakeRowsClaimed called without a result object"); } auto buffered = TakeBufferedRows(); unique_ptr rows; { DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; - auto taken = result->TakeChunks(check); + auto taken = result->TakeChunksClaimed(check); if (!buffered) { return taken; } @@ -495,29 +495,43 @@ T DuckDBPyResult::RunWithArrowSchema(const std::function return fun(arrow_schema); } +void DuckDBPyResult::BuildArrowSchemaClaimed(ArrowSchema &out) { + D_ASSERT(!Empty()); + result->BuildArrowSchemaClaimed(IdentifiersToStrings(ResultNames()), out); +} + duckdb::pyarrow::Table DuckDBPyResult::FetchArrowTable(const idx_t rows_per_batch, const bool to_polars) { if (Empty()) { throw InvalidInputException("There is no query result"); } RequireArrowFormat(*result); return RunWithArrowSchema( - [&](const ArrowSchema &schema) -> duckdb::pyarrow::Table { - auto pyarrow_schema = pyarrow::ToPyArrowSchema(schema); - vector> arrays; - { - DUCKDB_PY_ASSERT_GIL_HELD(); - nb::gil_scoped_release release; - arrays = result->TakeArrays(DuckDBPyConnection::CheckSignals); - } - nb::list batches; - for (auto &array : arrays) { - ArrowArray data; - array->MoveTo(data); - TransformDuckToArrowChunk(pyarrow_schema, data, batches); - } - return pyarrow::ToArrowTable(std::move(batches), pyarrow_schema); - }, - to_polars); + [&](const ArrowSchema &schema) -> duckdb::pyarrow::Table { return ConvertToArrowTable(schema); }, to_polars); +} + +duckdb::pyarrow::Table DuckDBPyResult::FetchArrowTable(ArrowSchema &schema) { + if (Empty()) { + throw InvalidInputException("There is no query result"); + } + RequireArrowFormat(*result); + return ConvertToArrowTable(schema); +} + +duckdb::pyarrow::Table DuckDBPyResult::ConvertToArrowTable(const ArrowSchema &schema) { + auto pyarrow_schema = pyarrow::ToPyArrowSchema(schema); + vector> arrays; + { + DUCKDB_PY_ASSERT_GIL_HELD(); + nb::gil_scoped_release release; + arrays = result->TakeArrays(DuckDBPyConnection::CheckSignals); + } + nb::list batches; + for (auto &array : arrays) { + ArrowArray data; + array->MoveTo(data); + TransformDuckToArrowChunk(pyarrow_schema, data, batches); + } + return pyarrow::ToArrowTable(std::move(batches), pyarrow_schema); } void DuckDBPyResult::CheckBatchSize(idx_t rows_per_batch) { @@ -652,10 +666,13 @@ struct ArrowArrayStreamGuard { } // namespace -ArrowArrayStream DuckDBPyResult::FetchArrowArrayStream(idx_t rows_per_batch) { +ArrowArrayStream DuckDBPyResult::FetchArrowArrayStream(idx_t rows_per_batch, ArrowSchema *prebuilt) { RequireArrowFormat(*result); auto state = make_uniq(); - if (names_override.empty()) { + if (prebuilt) { + state->cached_schema = *prebuilt; + prebuilt->release = nullptr; + } else if (names_override.empty()) { result->CopyArrowSchema(state->cached_schema); } else { result->BuildArrowSchema(IdentifiersToStrings(names_override), state->cached_schema); @@ -674,13 +691,22 @@ ArrowArrayStream DuckDBPyResult::FetchArrowArrayStream(idx_t rows_per_batch) { } duckdb::pyarrow::RecordBatchReader DuckDBPyResult::FetchRecordBatchReader(idx_t rows_per_batch) { + return ImportRecordBatchReader(nullptr, rows_per_batch); +} + +duckdb::pyarrow::RecordBatchReader DuckDBPyResult::FetchRecordBatchReader(ArrowSchema &schema, idx_t rows_per_batch) { + return ImportRecordBatchReader(&schema, rows_per_batch); +} + +duckdb::pyarrow::RecordBatchReader DuckDBPyResult::ImportRecordBatchReader(ArrowSchema *prebuilt, + idx_t rows_per_batch) { if (Empty()) { throw InvalidInputException("There is no query result"); } CheckBatchSize(rows_per_batch); auto pyarrow_lib_module = nb::module_::import_("pyarrow").attr("lib"); auto record_batch_reader_func = pyarrow_lib_module.attr("RecordBatchReader").attr("_import_from_c"); - ArrowArrayStreamGuard guard {FetchArrowArrayStream(rows_per_batch)}; + ArrowArrayStreamGuard guard {FetchArrowArrayStream(rows_per_batch, prebuilt)}; nb::object record_batch_reader = record_batch_reader_func((uint64_t)&guard.stream); // NOLINT return nb::cast(record_batch_reader); } @@ -701,7 +727,7 @@ nb::object DuckDBPyResult::FetchArrowCapsule(const idx_t rows_per_batch) { throw InvalidInputException("There is no query result"); } CheckBatchSize(rows_per_batch); - auto inner_stream = FetchArrowArrayStream(rows_per_batch); + auto inner_stream = FetchArrowArrayStream(rows_per_batch, nullptr); auto arrow_stream = new ArrowArrayStream(); *arrow_stream = inner_stream; return nb::capsule(arrow_stream, "arrow_array_stream", ArrowArrayStreamPyCapsuleDestructor); diff --git a/tests/fast/api/test_appender.py b/tests/fast/api/test_appender.py index cf2dfcb7..72c6473d 100644 --- a/tests/fast/api/test_appender.py +++ b/tests/fast/api/test_appender.py @@ -1757,6 +1757,83 @@ def test_another_thread_is_refused_during_a_flush_without_waiting(self, con, hel app.flush() assert table_rows(con) == [(1, "a")] + @pytest.mark.timeout(120) + def test_a_flush_that_collides_with_a_submitted_statement_is_refused_as_busy(self): + # The engine starts a statement as soon as another thread submits it, so a flush that lands right after the + # submission collides with a running statement: it is refused as busy, like one that lands later, instead + # of getting the engine's refusal or hanging the statement + code, output = run_isolated( + f""" + import collections + import threading + import time + import duckdb + from duckdb.sqltypes import BIGINT + con = duckdb.connect() + con.execute("CREATE TABLE t (k INTEGER, v VARCHAR)") + con.execute("CREATE TABLE log (k BIGINT)") + # Workers are what runs the statement from its submission on, so a one-CPU host must still get some + con.execute("SET threads = 4") + gates = [] + def hold(x): + gates[-1].wait(3) + return x + con.create_function("hold", hold, [BIGINT], BIGINT) + app = con.appender("t") + flushed = 0 + refusals = collections.Counter() + submitted = 0 + print("started", flush=True) + overall = time.monotonic() + 40 + while submitted < 20 and time.monotonic() < overall: + go, stop, release = threading.Event(), threading.Event(), threading.Event() + gates.append(release) + def flush_until_stopped(): + global flushed + # Started as the first thread enters its execute, so that call is the one under way; the + # statement is held through the first attempts and released while the attempts go on + go.wait(10) + attempts = 0 + try: + while not stop.is_set(): + app.append_row(flushed, "a") + try: + app.flush() + flushed += 1 + except duckdb.ResourceInUseException as error: + refusals["busy" if {BUSY!r} in str(error) else str(error)] += 1 + app.clear() + attempts += 1 + if attempts == 200: + release.set() + finally: + release.set() + thread = threading.Thread(target=flush_until_stopped) + thread.start() + go.set() + deadline = time.monotonic() + 1 + while time.monotonic() < deadline: + try: + con.execute("INSERT INTO log SELECT hold(i) FROM range(1) t(i)") + submitted += 1 + break + except duckdb.ResourceInUseException as error: + if {BUSY!r} not in str(error): + refusals["execute: " + str(error)] += 1 + stop.set() + thread.join(10) + print("submitted:", submitted) + print("refusals:", sorted(refusals)) + kept = con.execute("SELECT count(*) = ? FROM t", [flushed]).fetchall() + print("rows:", kept + con.execute("SELECT count(*) FROM log").fetchall()) + """, + timeout=90, + ) + assert code == 0, output + assert "submitted: 20" in output, output + assert "refusals: ['busy']" in output, output + assert "rows: [(True,), (20,)]" in output, output + @pytest.mark.parametrize("call", ["app.append_row(9, 'z')", "app.flush()", "app.clear()"]) def test_a_callback_during_its_flush_is_refused(self, call): code, output = run_isolated( diff --git a/tests/fast/api/test_execute_contract.py b/tests/fast/api/test_execute_contract.py index df91e510..11b66f06 100644 --- a/tests/fast/api/test_execute_contract.py +++ b/tests/fast/api/test_execute_contract.py @@ -1,3 +1,4 @@ +import contextlib import gc import importlib.util import re @@ -6,6 +7,7 @@ import textwrap import threading import time +from collections.abc import Iterator import pytest @@ -577,7 +579,10 @@ def inner(x): assert INSIDE_A_QUERY in output assert "usable after: [(42,)]" in output - def test_another_thread_is_refused_while_a_statement_runs(self, con): + def test_another_thread_is_refused_while_a_statement_runs(self): + # At the default thread count the engine runs the statement from its submission on, so the other thread + # may land right after the submission as well as during the settling + con = duckdb.connect() running, release = threading.Event(), threading.Event() def hold(x): @@ -602,9 +607,160 @@ def contend(): thread.start() assert con.execute("SELECT hold(i) FROM range(1) t(i)").fetchall() == [(0,)] thread.join(10) + con.close() assert len(outcome) == 1 assert BUSY in outcome[0] + @pytest.mark.timeout(120) + def test_a_second_execute_is_refused_instead_of_waiting(self): + # A collision refuses at once. An execute that passes the busy check while the other thread's statement is + # already submitted would wait for that statement instead, and a second execute never waits for one + code, output = run_isolated( + f""" + import collections + import threading + import time + import duckdb + from duckdb.sqltypes import BIGINT + con = duckdb.connect() + # Workers are what runs the statement from its submission on, so a one-CPU host must still get some + con.execute("SET threads = 4") + gates = [] + def hold(x): + gates[-1].wait(3) + return x + con.create_function("hold", hold, [BIGINT], BIGINT) + outcomes = collections.Counter() + submitted = 0 + print("started", flush=True) + overall = time.monotonic() + 40 + while submitted < 20 and time.monotonic() < overall: + go, release = threading.Event(), threading.Event() + gates.append(release) + def execute_two_hundred_times(): + # Started as the first thread enters its execute, so that call is the one under way; the + # statement is held until the attempts are over, so a call of this thread can replace the + # first thread's result only by waiting for it + go.wait(10) + try: + for attempt in range(200): + began = time.monotonic() + try: + con.execute("SELECT 7") + except duckdb.ResourceInUseException as error: + if {BUSY!r} not in str(error): + outcomes[str(error)] += 1 + if time.monotonic() - began > 1: + outcomes["waited"] += 1 + finally: + release.set() + thread = threading.Thread(target=execute_two_hundred_times) + thread.start() + go.set() + deadline = time.monotonic() + 1 + got_through = False + while time.monotonic() < deadline: + try: + con.execute("SELECT hold(i) + 10 FROM range(1) t(i)") + submitted += 1 + got_through = True + break + except duckdb.ResourceInUseException as error: + if {BUSY!r} not in str(error): + outcomes["execute: " + str(error)] += 1 + thread.join(10) + if got_through: + rows = con.fetchall() + if rows != [(10,)]: + outcomes[f"first thread read {{rows}}"] += 1 + print("submitted:", submitted) + print("outcomes:", sorted(outcomes)) + """, + timeout=90, + ) + assert code == 0, output + assert "submitted: 20" in output, output + assert "outcomes: []" in output, output + + def test_a_second_execute_during_parameter_conversion_is_refused_at_once(self): + # The first thread holds the connection mutex and no claim while its parameters are read, so a second + # execute reaches the mutex; it is refused there, and does not wait for the first call and replace its result + con = duckdb.connect() + inside, done = threading.Event(), threading.Event() + timed_out = [] + + class Params: + def __len__(self) -> int: + return 1 + + def __iter__(self) -> Iterator[int]: + inside.set() + # The second thread ends the wait once its call returned; a call that waits for this one never does + timed_out.append(not done.wait(5)) + return iter([1]) + + outcome = [] + + def second(): + assert inside.wait(10) + try: + con.execute("SELECT 2") + outcome.append("ran") + except duckdb.ResourceInUseException as error: + outcome.append(str(error)) + finally: + done.set() + + thread = threading.Thread(target=second) + thread.start() + try: + con.execute("SELECT ?::INTEGER", Params()) + rows = con.fetchall() + finally: + thread.join(10) + con.close() + assert timed_out == [False], outcome + assert len(outcome) == 1, outcome + assert BUSY in outcome[0], outcome + assert rows == [(1,)] + + def test_a_thread_reading_the_description_does_not_make_execute_fail(self): + # A reader holds the connection mutex while it waits for the GIL back; an execute on another thread waits for + # such a holder as every call did, and is refused only behind another call that submits a statement + con = duckdb.connect() + con.execute("SELECT 1") + stop = threading.Event() + + def read_until_stopped(): + while not stop.is_set(): + with contextlib.suppress(duckdb.ResourceInUseException): + con.description # noqa: B018 + time.sleep(0.0005) + + def execute_once(): + try: + con.execute("SELECT 1") + except duckdb.ResourceInUseException: + return False + return True + + thread = threading.Thread(target=read_until_stopped) + thread.start() + ran, refused = 0, 0 + deadline = time.monotonic() + 1 + try: + while time.monotonic() < deadline: + if execute_once(): + ran += 1 + else: + refused += 1 + finally: + stop.set() + thread.join(10) + con.close() + assert refused == 0, (ran, refused) + assert ran > 0 + @pytest.mark.timeout(60) def test_progress_and_interrupt_reach_a_running_statement(self, con): seen = [] From 8e73d3f1aaafd38cdac51f9a08975031a2db95be Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Tue, 10 Nov 2026 23:08:13 +0000 Subject: [PATCH 2/6] Refuse a submitting call on arrival while another thread's submitting call is in flight, and pin each closed gap with a test --- src/engine/connection.cpp | 3 + src/engine/result.cpp | 2 + .../duckdb_python/engine/connection.hpp | 4 + .../pyconnection/pyconnection.hpp | 61 +++---- src/include/duckdb_python/pyresult.hpp | 5 +- src/pyconnection.cpp | 33 ++-- src/pyresult.cpp | 5 - tests/fast/api/test_appender.py | 78 +++++++++ tests/fast/api/test_execute_contract.py | 165 +++++++++++++++++- 9 files changed, 286 insertions(+), 70 deletions(-) diff --git a/src/engine/connection.cpp b/src/engine/connection.cpp index 1f98f194..d9f8cf96 100644 --- a/src/engine/connection.cpp +++ b/src/engine/connection.cpp @@ -346,6 +346,7 @@ ConnectionClaim Session::Claim() { shared_ptr Session::SubmitClaimed(unique_ptr statement, identifier_map_t &values, const Format &format) { + D_ASSERT(state->Claimed()); auto written_type = statement->type; registry->CountStatement(written_type); auto handle = context->Submit(std::move(statement), values, QueryParameters(format.EngineFormat())); @@ -357,6 +358,7 @@ shared_ptr Session::SubmitClaimed(unique_ptr statement, shared_ptr Session::SubmitClaimed(Prepared &prepared, identifier_map_t &values, const Format &format) { + D_ASSERT(state->Claimed()); registry->CountStatement(prepared.statement->GetStatementType()); unique_ptr handle; { @@ -371,6 +373,7 @@ shared_ptr Session::SubmitClaimed(Prepared &prepared, identifier_map_t Session::SubmitClaimed(unique_ptr collection, const vector &names, const Format &format) { + D_ASSERT(state->Claimed()); // The binder rejects duplicate column names; the caller keeps the originals auto deduplicated_names = names; QueryResult::DeduplicateColumns(deduplicated_names); diff --git a/src/engine/result.cpp b/src/engine/result.cpp index aae5de4a..5137ae2c 100644 --- a/src/engine/result.cpp +++ b/src/engine/result.cpp @@ -265,6 +265,7 @@ void Result::RetainClaimed(const InterruptCheck &check) { if (state != State::PENDING) { return; } + D_ASSERT(connection_state->Claimed()); try { Complete(*handle, check); CaptureMetadata(); @@ -280,6 +281,7 @@ void Result::SettleClaimed(const InterruptCheck &check) { if (state != State::PENDING) { return; } + D_ASSERT(connection_state->Claimed()); if (CompletesBeforeReturning()) { RetainClaimed(check); return; diff --git a/src/include/duckdb_python/engine/connection.hpp b/src/include/duckdb_python/engine/connection.hpp index be4c9687..a9a28226 100644 --- a/src/include/duckdb_python/engine/connection.hpp +++ b/src/include/duckdb_python/engine/connection.hpp @@ -86,6 +86,10 @@ class ConnectionState : public ClientContextState { bool AdmittedDuringExit() const { return admitted_during_exit; } + //! Whether a call holds the connection's claim right now + bool Claimed() const { + return busy; + } private: friend class ConnectionClaim; diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 03a85cf4..06ecb324 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -7,6 +7,8 @@ //===----------------------------------------------------------------------===// #pragma once + +#include #include "duckdb.hpp" #include "duckdb_python/nb/casters.hpp" #include "duckdb_python/import_cache/python_import_cache.hpp" @@ -267,13 +269,15 @@ struct DuckDBPyConnection : public std::enable_shared_from_this mutex_held_by_submitting_call {false}; - //! How many times a submitting call has taken py_connection_lock - std::atomic submitting_entries {0}; + //! The thread whose call that submits a statement (IfHeld::REFUSE) is in flight, or no thread + std::atomic submitting_thread {std::thread::id()}; //! The database's store of file-like objects for reading, kept once this connection has found it shared_ptr object_store; case_insensitive_map_t> registered_functions; diff --git a/src/include/duckdb_python/pyresult.hpp b/src/include/duckdb_python/pyresult.hpp index 28188411..bb1e1da6 100644 --- a/src/include/duckdb_python/pyresult.hpp +++ b/src/include/duckdb_python/pyresult.hpp @@ -76,12 +76,9 @@ struct DuckDBPyResult { duckdb::pyarrow::Table FetchArrowTable(idx_t rows_per_batch, bool to_polars); duckdb::pyarrow::RecordBatchReader FetchRecordBatchReader(idx_t rows_per_batch = DEFAULT_ARROW_BATCH_SIZE); - //! The two fetches above over a schema built with BuildArrowSchemaClaimed, which they take over + //! The two fetches above over a schema the caller built for the result's names, which they take over duckdb::pyarrow::Table FetchArrowTable(ArrowSchema &schema); duckdb::pyarrow::RecordBatchReader FetchRecordBatchReader(ArrowSchema &schema, idx_t rows_per_batch); - //! The result's Arrow schema under the names the Python layer reports, for a caller that holds the connection's - //! claim; out is the caller's to release until a fetch takes it over - void BuildArrowSchemaClaimed(ArrowSchema &out); nb::object FetchArrowCapsule(idx_t rows_per_batch = DEFAULT_ARROW_BATCH_SIZE); static void CheckBatchSize(idx_t rows_per_batch); diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index 12b1ce65..427dc034 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -1508,17 +1508,6 @@ nb::dict DuckDBPyConnection::FetchTF() { return result.FetchTF(); } -//! Releases an Arrow schema that was never handed over, after a throw before the import into pyarrow, which takes the -//! schema over even when it fails -struct ArrowSchemaHolder { - ArrowSchema schema {}; - ~ArrowSchemaHolder() { - if (schema.release) { - schema.release(&schema); - } - } -}; - nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { ConnectionLockGuard conn_lock(*this, IfHeld::REFUSE); // Asked before the result, so a closed connection reports itself and not the slot its close emptied @@ -1540,8 +1529,10 @@ nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { // Imported before the engine checks, so no Python code runs between what they find and the take nb::module_::import_("pyarrow"); auto held = con.GetResult().EngineResult(); - unique_ptr converted; - ArrowSchemaHolder schema; + shared_ptr arrow; + vector names; + // Released by the wrapper when a throw before the import into pyarrow leaves it unused; the import takes it over + ArrowSchemaWrapper schema; { // One claim from the checks through the take of the rows, the statement over them, its settle or retain and // the schema build, so a refusal comes before the rows are taken and never loses them. The claim is a flag @@ -1560,12 +1551,14 @@ nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { } } - vector names; auto rows = con.GetResult().TakeRowsClaimed(names); if (!rows) { return nb::none(); } - shared_ptr arrow; + vector schema_names; + for (auto &name : names) { + schema_names.push_back(name.GetIdentifierName()); + } { DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; @@ -1575,18 +1568,18 @@ nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { } else { arrow->RetainClaimed(DuckDBPyConnection::CheckSignals); } + // The converter reads the connection's settings and its catalog's types, so the schema is built here + arrow->BuildArrowSchemaClaimed(schema_names, schema.arrow_schema); } - converted = make_uniq(std::move(arrow), std::move(names)); - // The converter reads the connection's settings and its catalog's types, so the schema is built here - converted->BuildArrowSchemaClaimed(schema.schema); } // The rows are retained, or the stream is settled, so the conversion into pyarrow holds the connection no longer // and refuses no other call conn_lock.DropSubmittingMark(); + DuckDBPyResult converted(std::move(arrow), std::move(names)); if (reader) { - return nb::cast(converted->FetchRecordBatchReader(schema.schema, batch_size)); + return nb::cast(converted.FetchRecordBatchReader(schema.arrow_schema, batch_size)); } - return nb::cast(converted->FetchArrowTable(schema.schema)); + return nb::cast(converted.FetchArrowTable(schema.arrow_schema)); } identifier_map_t TransformPyConfigDict(const nb::dict &py_config_dict) { diff --git a/src/pyresult.cpp b/src/pyresult.cpp index 74b8c5ab..666d065e 100644 --- a/src/pyresult.cpp +++ b/src/pyresult.cpp @@ -495,11 +495,6 @@ T DuckDBPyResult::RunWithArrowSchema(const std::function return fun(arrow_schema); } -void DuckDBPyResult::BuildArrowSchemaClaimed(ArrowSchema &out) { - D_ASSERT(!Empty()); - result->BuildArrowSchemaClaimed(IdentifiersToStrings(ResultNames()), out); -} - duckdb::pyarrow::Table DuckDBPyResult::FetchArrowTable(const idx_t rows_per_batch, const bool to_polars) { if (Empty()) { throw InvalidInputException("There is no query result"); diff --git a/tests/fast/api/test_appender.py b/tests/fast/api/test_appender.py index 72c6473d..e4e57b8e 100644 --- a/tests/fast/api/test_appender.py +++ b/tests/fast/api/test_appender.py @@ -1834,6 +1834,84 @@ def flush_until_stopped(): assert "refusals: ['busy']" in output, output assert "rows: [(True,), (20,)]" in output, output + @pytest.mark.timeout(120) + def test_a_flush_that_collides_with_a_submitted_executemany_is_refused_as_busy(self): + # As above, with the statement submitted by executemany, which runs it once per parameter set + code, output = run_isolated( + f""" + import collections + import threading + import time + import duckdb + from duckdb.sqltypes import BIGINT + con = duckdb.connect() + con.execute("CREATE TABLE t (k INTEGER, v VARCHAR)") + con.execute("CREATE TABLE log (k BIGINT)") + # Workers are what runs the statement from its submission on, so a one-CPU host must still get some + con.execute("SET threads = 4") + gates = [] + def hold(x): + gates[-1].wait(3) + return x + con.create_function("hold", hold, [BIGINT], BIGINT) + app = con.appender("t") + flushed = 0 + refusals = collections.Counter() + submitted = 0 + print("started", flush=True) + overall = time.monotonic() + 40 + while submitted < 20 and time.monotonic() < overall: + go, stop, release = threading.Event(), threading.Event(), threading.Event() + gates.append(release) + def flush_until_stopped(): + global flushed + # Started as the first thread enters its execute, so that call is the one under way; the + # statement is held through the first attempts and released while the attempts go on + go.wait(10) + attempts = 0 + try: + while not stop.is_set(): + app.append_row(flushed, "a") + try: + app.flush() + flushed += 1 + except duckdb.ResourceInUseException as error: + refusals["busy" if {BUSY!r} in str(error) else str(error)] += 1 + app.clear() + attempts += 1 + if attempts == 200: + release.set() + finally: + release.set() + thread = threading.Thread(target=flush_until_stopped) + thread.start() + go.set() + deadline = time.monotonic() + 1 + while time.monotonic() < deadline: + try: + con.executemany("INSERT INTO log SELECT hold(i + ?::BIGINT) FROM range(1) t(i)", [[1], [2]]) + submitted += 1 + break + except duckdb.ResourceInUseException as error: + if {BUSY!r} not in str(error): + refusals["execute: " + str(error)] += 1 + stop.set() + thread.join(10) + print("submitted:", submitted) + print("refusals:", sorted(refusals)) + kept = con.execute("SELECT count(*) = ? FROM t", [flushed]).fetchall() + # An executemany refused between its two parameter sets, which is allowed, has inserted its first row + # already and is retried in full, so the log holds at least two rows per submission + logged = con.execute("SELECT count(*) >= 40 FROM log").fetchall() + print("rows:", kept + logged) + """, + timeout=90, + ) + assert code == 0, output + assert "submitted: 20" in output, output + assert "refusals: ['busy']" in output, output + assert "rows: [(True,), (True,)]" in output, output + @pytest.mark.parametrize("call", ["app.append_row(9, 'z')", "app.flush()", "app.clear()"]) def test_a_callback_during_its_flush_is_refused(self, call): code, output = run_isolated( diff --git a/tests/fast/api/test_execute_contract.py b/tests/fast/api/test_execute_contract.py index 11b66f06..b3ecb253 100644 --- a/tests/fast/api/test_execute_contract.py +++ b/tests/fast/api/test_execute_contract.py @@ -441,7 +441,7 @@ def test_a_lazy_polars_frame_scanned_on_its_own_connection_errors(self): ) assert code == 0, output assert "refused:" in output - assert re.search(REFUSED, output) + assert BUSY in output, output assert "usable after: [(42,)]" in output def test_a_udf_that_reenters_its_connection_errors(self): @@ -580,9 +580,11 @@ def inner(x): assert "usable after: [(42,)]" in output def test_another_thread_is_refused_while_a_statement_runs(self): - # At the default thread count the engine runs the statement from its submission on, so the other thread - # may land right after the submission as well as during the settling + # With worker threads the engine runs the statement from its submission on, so the other thread may land + # right after the submission as well as during the settling. It contends through a relation, which takes + # the connection's claim and not the mutex, so it sees the claim and nothing else con = duckdb.connect() + con.execute("SET threads = 4") running, release = threading.Event(), threading.Event() def hold(x): @@ -596,7 +598,7 @@ def hold(x): def contend(): running.wait(10) try: - con.execute("SELECT 1") + con.sql("SELECT 1") outcome.append("ran") except duckdb.ResourceInUseException as error: outcome.append(str(error)) @@ -682,10 +684,65 @@ def execute_two_hundred_times(): assert "submitted: 20" in output, output assert "outcomes: []" in output, output + @pytest.mark.timeout(120) + def test_a_refused_arrow_fetch_keeps_its_rows(self): + # The DB-API Arrow fetch takes the rows out of the result and runs a statement over them; a second thread + # that takes the connection's claim between the two would make the fetch fail after the take, and a retry + # would then find no rows. The fetch holds one claim across both, so a refusal always comes before the take + code, output = run_isolated( + f""" + import collections + import sys + import threading + import time + import duckdb + con = duckdb.connect() + con.execute("SET threads = 4") + # A short switch interval hands the GIL to the other thread often, as a busy program does + sys.setswitchinterval(0.0001) + stop = threading.Event() + def bind_until_stopped(): + while not stop.is_set(): + try: + con.sql("SELECT 1") + except duckdb.ResourceInUseException: + pass + thread = threading.Thread(target=bind_until_stopped) + thread.start() + counts = collections.Counter() + print("started", flush=True) + deadline = time.monotonic() + 8 + while time.monotonic() < deadline and counts["retried"] < 50: + try: + con.execute("SELECT * FROM range(5) t(i)") + except duckdb.ResourceInUseException: + continue + for attempt in range(1000): + try: + table = con.to_arrow_table() + break + except duckdb.ResourceInUseException as error: + if {BUSY!r} not in str(error): + counts["refused: " + str(error)] += 1 + rows = None if table is None else table.num_rows + if attempt: + counts["retried"] += 1 + if rows != 5: + counts[f"rows {{rows}}" + (" after a refusal" if attempt else "")] += 1 + stop.set() + thread.join(10) + print("counts:", sorted(counts)) + """, + timeout=90, + ) + assert code == 0, output + assert "counts: ['retried']" in output or "counts: []" in output, output + def test_a_second_execute_during_parameter_conversion_is_refused_at_once(self): # The first thread holds the connection mutex and no claim while its parameters are read, so a second # execute reaches the mutex; it is refused there, and does not wait for the first call and replace its result con = duckdb.connect() + con.execute("SET threads = 4") inside, done = threading.Event(), threading.Event() timed_out = [] @@ -728,6 +785,7 @@ def test_a_thread_reading_the_description_does_not_make_execute_fail(self): # A reader holds the connection mutex while it waits for the GIL back; an execute on another thread waits for # such a holder as every call did, and is refused only behind another call that submits a statement con = duckdb.connect() + con.execute("SET threads = 4") con.execute("SELECT 1") stop = threading.Event() @@ -761,6 +819,105 @@ def execute_once(): assert refused == 0, (ran, refused) assert ran > 0 + def test_a_submitting_call_beside_a_reader_is_refused_at_once(self): + # A reader holding the mutex puts two submitting calls in line for it; the refusal of the second is decided + # when it arrives, so it never waits for the first one's statement and is then refused + con = duckdb.connect() + con.execute("SET threads = 4") + # The held statement lasts well over any scheduling stall, so only a wait for it crosses the bound below + con.create_function("hold", lambda x: time.sleep(0.3) or x, [BIGINT], BIGINT) + stop = threading.Event() + waited = [] + + def read_until_stopped(): + while not stop.is_set(): + with contextlib.suppress(duckdb.ResourceInUseException): + con.description # noqa: B018 + + def run_held_statements(): + while not stop.is_set(): + with contextlib.suppress(duckdb.ResourceInUseException): + con.execute("SELECT hold(i) FROM range(1) t(i)") + + def execute_once(): + began = time.monotonic() + with contextlib.suppress(duckdb.ResourceInUseException): + con.execute("SELECT 7") + return time.monotonic() - began + + threads = [threading.Thread(target=read_until_stopped), threading.Thread(target=run_held_statements)] + for thread in threads: + thread.start() + deadline = time.monotonic() + 4 + try: + while time.monotonic() < deadline: + took = execute_once() + if took > 0.15: + waited.append(took) + finally: + stop.set() + for thread in threads: + thread.join(10) + con.close() + assert waited == [] + + def test_other_calls_run_while_a_result_is_converted_to_arrow(self): + # The Arrow fetch holds the connection only until the rows are taken and the schema is built; the conversion + # into pyarrow that follows holds the mutex and no claim, so a relation runs during it and an execute waits + # for it instead of being refused. Small batches make the conversion the long part of the fetch + con = duckdb.connect() + con.execute("SET threads = 4") + con.execute("CREATE TABLE t (k INTEGER)") + converting = threading.Event() + marks = {} + attempts = [] + outcome = {} + + def attempt(call): + try: + call() + except duckdb.ResourceInUseException as error: + return str(error) + return "ran" + + def second(): + assert converting.wait(10) + deadline = time.monotonic() + 20 + while time.monotonic() < deadline and "done" not in marks: + began = time.monotonic() + kind = attempt(lambda: con.table("t")) + attempts.append((began, kind)) + # A relation admitted after the take was refused ran during the conversion; an execute then too + if kind == "ran" and any(k != "ran" for _, k in attempts): + outcome["execute during the conversion"] = attempt(lambda: con.execute("SELECT 1")) + break + + def fetch_or_none(): + try: + return con.to_arrow_table(200) + except duckdb.ResourceInUseException: + return None + + thread = threading.Thread(target=second) + thread.start() + try: + con.execute("SELECT i, i::VARCHAR AS s FROM range(1000000) t(i)") + converting.set() + table = None + while table is None: + table = fetch_or_none() + marks["done"] = time.monotonic() + finally: + marks.setdefault("done", time.monotonic()) + thread.join(30) + con.close() + assert table.num_rows == 1000000 + refused = [t for t, k in attempts if k != "ran"] + assert refused, attempts[:5] + ran_after_take = [t for t, k in attempts if k == "ran" and t > max(refused) and t < marks["done"]] + assert ran_after_take, (len(attempts), max(refused), marks["done"]) + assert outcome == {"execute during the conversion": "ran"}, outcome + @pytest.mark.timeout(60) def test_progress_and_interrupt_reach_a_running_statement(self, con): seen = [] From 27e7efbf930f91c62880f05f37e66a0928bf51e2 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Wed, 11 Nov 2026 00:13:45 +0000 Subject: [PATCH 3/6] Let the executemany collision test retry past the refusal between parameter sets --- tests/fast/api/test_appender.py | 32 +++++++++++++++----------------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/tests/fast/api/test_appender.py b/tests/fast/api/test_appender.py index e4e57b8e..b94000fe 100644 --- a/tests/fast/api/test_appender.py +++ b/tests/fast/api/test_appender.py @@ -1836,7 +1836,9 @@ def flush_until_stopped(): @pytest.mark.timeout(120) def test_a_flush_that_collides_with_a_submitted_executemany_is_refused_as_busy(self): - # As above, with the statement submitted by executemany, which runs it once per parameter set + # As above, with the statement submitted by executemany, which runs it once per parameter set. Between two + # sets the connection is free, so a flush landing there runs and the executemany is refused at its next set; + # that is allowed, and the main thread retries it, so the second thread stops after its attempts code, output = run_isolated( f""" import collections @@ -1859,18 +1861,17 @@ def hold(x): refusals = collections.Counter() submitted = 0 print("started", flush=True) - overall = time.monotonic() + 40 + overall = time.monotonic() + 60 while submitted < 20 and time.monotonic() < overall: - go, stop, release = threading.Event(), threading.Event(), threading.Event() + go, release = threading.Event(), threading.Event() gates.append(release) - def flush_until_stopped(): + def flush_two_hundred_times(): global flushed - # Started as the first thread enters its execute, so that call is the one under way; the + # Started as the first thread enters its executemany, so that call is the one under way; the # statement is held through the first attempts and released while the attempts go on go.wait(10) - attempts = 0 try: - while not stop.is_set(): + for attempt in range(200): app.append_row(flushed, "a") try: app.flush() @@ -1878,31 +1879,28 @@ def flush_until_stopped(): except duckdb.ResourceInUseException as error: refusals["busy" if {BUSY!r} in str(error) else str(error)] += 1 app.clear() - attempts += 1 - if attempts == 200: + if attempt == 100: release.set() finally: release.set() - thread = threading.Thread(target=flush_until_stopped) + thread = threading.Thread(target=flush_two_hundred_times) thread.start() go.set() - deadline = time.monotonic() + 1 - while time.monotonic() < deadline: + while time.monotonic() < overall: try: con.executemany("INSERT INTO log SELECT hold(i + ?::BIGINT) FROM range(1) t(i)", [[1], [2]]) submitted += 1 break except duckdb.ResourceInUseException as error: if {BUSY!r} not in str(error): - refusals["execute: " + str(error)] += 1 - stop.set() + refusals["executemany: " + str(error)] += 1 thread.join(10) print("submitted:", submitted) print("refusals:", sorted(refusals)) kept = con.execute("SELECT count(*) = ? FROM t", [flushed]).fetchall() - # An executemany refused between its two parameter sets, which is allowed, has inserted its first row - # already and is retried in full, so the log holds at least two rows per submission - logged = con.execute("SELECT count(*) >= 40 FROM log").fetchall() + # An executemany refused between its two parameter sets has written its first row already and is retried + # in full, so the log holds at least two rows per submission + logged = con.execute("SELECT count(*) >= 2 * ? FROM log", [submitted]).fetchall() print("rows:", kept + logged) """, timeout=90, From ad67684bcd47393d981c4e841ed7f99d5d427f44 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Wed, 11 Nov 2026 01:13:08 +0000 Subject: [PATCH 4/6] Register the DB-API result inside execute's claim, mark the Arrow tests as needing pyarrow, bound the retry loops --- .../pyconnection/pyconnection.hpp | 6 +- src/pyconnection.cpp | 27 ++- src/pyresult.cpp | 8 +- tests/fast/api/test_appender.py | 6 + tests/fast/api/test_execute_contract.py | 193 ++++++++++++++---- 5 files changed, 184 insertions(+), 56 deletions(-) diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 06ecb324..01444999 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -282,7 +282,9 @@ struct DuckDBPyConnection : public std::enable_shared_from_this Execute(const nb::object &query, nb::object params = nb::list()); + //! Makes the result the connection's DB-API result, under the caller's claim; false when there is none + bool RegisterResult(shared_ptr res); std::shared_ptr ExecuteFromString(const string &query); std::shared_ptr Append(const string &name, const PandasDataFrame &value, bool by_name); diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index 427dc034..40623c93 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -882,6 +882,7 @@ std::shared_ptr DuckDBPyConnection::Execute(const nb::object RefuseParameters(statements); } shared_ptr res; + bool registered = false; if (last_statement) { if (params.is_none()) { params = nb::list(); @@ -891,27 +892,38 @@ std::shared_ptr DuckDBPyConnection::Execute(const nb::object DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; unique_lock lock(py_connection_lock); - // One claim from submission to settling, so no other call finds the statement pending + // One claim from submission to settling and on to the registration of the result, so no other call finds + // the statement pending, and one that takes the claim next meets the registered result and gets the + // client's refusal auto claim = session.Claim(); res = SubmitAndDrive(statements, parameters_given, false, [&]() { return session.SubmitClaimed(std::move(last_statement), named_values, engine::Format::Chunks()); }); - res = statements.Choose(std::move(res)); + registered = RegisterResult(statements.Choose(std::move(res))); } else { DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; unique_lock lock(py_connection_lock); - res = statements.Choose(nullptr); + registered = RegisterResult(statements.Choose(nullptr)); } - if (!res) { + if (!registered) { // TODO: should we throw? return nullptr; } - auto py_result = std::make_shared(std::move(res)); - con.SetResult(std::make_unique(std::move(py_result))); return shared_from_this(); } +bool DuckDBPyConnection::RegisterResult(shared_ptr res) { + DUCKDB_PY_ASSERT_GIL_RELEASED(); + if (!res) { + return false; + } + // The result the call superseded is gone already, so nothing that needs the GIL is destroyed here + D_ASSERT(!con.HasResult()); + con.SetResult(make_uniq(std::make_shared(std::move(res)))); + return true; +} + std::unique_ptr DuckDBPyConnection::CreateAppender(const string &table, const std::optional &columns, const string &action, @@ -1573,7 +1585,8 @@ nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { } } // The rows are retained, or the stream is settled, so the conversion into pyarrow holds the connection no longer - // and refuses no other call + // and refuses no submitting call; a settled stream stays the connection's open result until it is read out, as a + // relation's reader does conn_lock.DropSubmittingMark(); DuckDBPyResult converted(std::move(arrow), std::move(names)); if (reader) { diff --git a/src/pyresult.cpp b/src/pyresult.cpp index 666d065e..3ff6f9d6 100644 --- a/src/pyresult.cpp +++ b/src/pyresult.cpp @@ -667,10 +667,12 @@ ArrowArrayStream DuckDBPyResult::FetchArrowArrayStream(idx_t rows_per_batch, Arr if (prebuilt) { state->cached_schema = *prebuilt; prebuilt->release = nullptr; - } else if (names_override.empty()) { - result->CopyArrowSchema(state->cached_schema); } else { - result->BuildArrowSchema(IdentifiersToStrings(names_override), state->cached_schema); + // A result with its own names comes with a prebuilt schema, see ResultAsArrow + if (!names_override.empty()) { + throw InternalException("Arrow stream over a result with its own names and no schema built for them"); + } + result->CopyArrowSchema(state->cached_schema); } state->result = std::move(result); current_chunk.reset(); diff --git a/tests/fast/api/test_appender.py b/tests/fast/api/test_appender.py index b94000fe..3005f633 100644 --- a/tests/fast/api/test_appender.py +++ b/tests/fast/api/test_appender.py @@ -1803,6 +1803,9 @@ def flush_until_stopped(): except duckdb.ResourceInUseException as error: refusals["busy" if {BUSY!r} in str(error) else str(error)] += 1 app.clear() + except Exception as error: + refusals["second thread: " + type(error).__name__] += 1 + return attempts += 1 if attempts == 200: release.set() @@ -1879,6 +1882,9 @@ def flush_two_hundred_times(): except duckdb.ResourceInUseException as error: refusals["busy" if {BUSY!r} in str(error) else str(error)] += 1 app.clear() + except Exception as error: + refusals["second thread: " + type(error).__name__] += 1 + return if attempt == 100: release.set() finally: diff --git a/tests/fast/api/test_execute_contract.py b/tests/fast/api/test_execute_contract.py index b3ecb253..f0fe1e07 100644 --- a/tests/fast/api/test_execute_contract.py +++ b/tests/fast/api/test_execute_contract.py @@ -596,11 +596,11 @@ def hold(x): outcome = [] def contend(): - running.wait(10) try: + assert running.wait(10), "the statement never reached its UDF" con.sql("SELECT 1") outcome.append("ran") - except duckdb.ResourceInUseException as error: + except Exception as error: outcome.append(str(error)) finally: release.set() @@ -652,6 +652,8 @@ def execute_two_hundred_times(): except duckdb.ResourceInUseException as error: if {BUSY!r} not in str(error): outcomes[str(error)] += 1 + except Exception as error: + outcomes["second thread: " + type(error).__name__] += 1 if time.monotonic() - began > 1: outcomes["waited"] += 1 finally: @@ -684,6 +686,7 @@ def execute_two_hundred_times(): assert "submitted: 20" in output, output assert "outcomes: []" in output, output + @needs_pyarrow @pytest.mark.timeout(120) def test_a_refused_arrow_fetch_keeps_its_rows(self): # The DB-API Arrow fetch takes the rows out of the result and runs a statement over them; a second thread @@ -696,47 +699,62 @@ def test_a_refused_arrow_fetch_keeps_its_rows(self): import threading import time import duckdb + import pyarrow # noqa: F401 con = duckdb.connect() con.execute("SET threads = 4") # A short switch interval hands the GIL to the other thread often, as a busy program does sys.setswitchinterval(0.0001) stop = threading.Event() + counts = collections.Counter() def bind_until_stopped(): while not stop.is_set(): try: con.sql("SELECT 1") except duckdb.ResourceInUseException: pass - thread = threading.Thread(target=bind_until_stopped) + except Exception as error: + counts["binder: " + type(error).__name__] += 1 + return + thread = threading.Thread(target=bind_until_stopped, daemon=True) thread.start() - counts = collections.Counter() + fetches = retried = 0 print("started", flush=True) - deadline = time.monotonic() + 8 - while time.monotonic() < deadline and counts["retried"] < 50: - try: - con.execute("SELECT * FROM range(5) t(i)") - except duckdb.ResourceInUseException: - continue - for attempt in range(1000): + try: + deadline = time.monotonic() + 8 + while time.monotonic() < deadline and retried < 50: try: - table = con.to_arrow_table() - break - except duckdb.ResourceInUseException as error: - if {BUSY!r} not in str(error): - counts["refused: " + str(error)] += 1 - rows = None if table is None else table.num_rows - if attempt: - counts["retried"] += 1 - if rows != 5: - counts[f"rows {{rows}}" + (" after a refusal" if attempt else "")] += 1 - stop.set() - thread.join(10) + con.execute("SELECT * FROM range(5) t(i)") + except duckdb.ResourceInUseException: + continue + table = None + for attempt in range(1000): + try: + table = con.to_arrow_table() + break + except duckdb.ResourceInUseException as error: + if {BUSY!r} not in str(error): + counts["refused: " + str(error)] += 1 + fetches += 1 + if attempt: + retried += 1 + if table is None: + counts["no table" + (" after a refusal" if attempt else "")] += 1 + elif table.num_rows != 5: + counts[f"rows {{table.num_rows}}" + (" after a refusal" if attempt else "")] += 1 + finally: + stop.set() + thread.join(10) + print("fetches:", fetches, "retried:", retried) print("counts:", sorted(counts)) """, timeout=90, ) assert code == 0, output - assert "counts: ['retried']" in output or "counts: []" in output, output + assert "counts: []" in output, output + fetches, retried = (int(n) for n in re.search(r"fetches: (\d+) retried: (\d+)", output).groups()) + assert fetches > 0, output + if retried == 0: + pytest.skip("inconclusive: no fetch was refused by the binding thread") def test_a_second_execute_during_parameter_conversion_is_refused_at_once(self): # The first thread holds the connection mutex and no claim while its parameters are read, so a second @@ -781,6 +799,70 @@ def second(): assert BUSY in outcome[0], outcome assert rows == [(1,)] + def test_a_nested_execute_from_the_submitting_thread_still_enters(self): + # The thread whose submitting call is in flight may submit again from inside it, as the recursive mutex + # admits it: a parameter object that runs a statement reads its row, and the outer call returns its own + con = duckdb.connect() + con.execute("SET threads = 4") + inner = [] + + class Params: + def __len__(self) -> int: + return 1 + + def __iter__(self) -> Iterator[int]: + inner.append(con.execute("SELECT 99").fetchall()) + return iter([1]) + + try: + rows = con.execute("SELECT ?::INTEGER", Params()).fetchall() + finally: + con.close() + assert inner == [[(99,)]] + assert rows == [(1,)] + + def test_a_reader_is_refused_while_a_statement_runs(self): + # A call that waits for the mutex when a reader holds it is still refused outright while a statement holds + # the connection: the busy check comes before the wait + con = duckdb.connect() + con.execute("SET threads = 4") + running, release = threading.Event(), threading.Event() + + def hold(x): + running.set() + release.wait(10) + return x + + con.create_function("hold", hold, [BIGINT], BIGINT) + outcome = {} + + def attempt(name, call): + try: + call() + outcome[name] = "ran" + except duckdb.ResourceInUseException as error: + outcome[name] = "busy" if BUSY in str(error) else str(error) + + def contend(): + try: + assert running.wait(10), "the statement never reached its UDF" + attempt("description", lambda: con.description) + attempt("fetchall", con.fetchall) + except Exception as error: + outcome["contender failed"] = repr(error) + finally: + release.set() + + thread = threading.Thread(target=contend) + thread.start() + try: + rows = con.execute("SELECT hold(i) FROM range(1) t(i)").fetchall() + finally: + thread.join(10) + con.close() + assert rows == [(0,)] + assert outcome == {"description": "busy", "fetchall": "busy"}, outcome + def test_a_thread_reading_the_description_does_not_make_execute_fail(self): # A reader holds the connection mutex while it waits for the GIL back; an execute on another thread waits for # such a holder as every call did, and is refused only behind another call that submits a statement @@ -789,11 +871,16 @@ def test_a_thread_reading_the_description_does_not_make_execute_fail(self): con.execute("SELECT 1") stop = threading.Event() + failures = [] + def read_until_stopped(): - while not stop.is_set(): - with contextlib.suppress(duckdb.ResourceInUseException): - con.description # noqa: B018 - time.sleep(0.0005) + try: + while not stop.is_set(): + with contextlib.suppress(duckdb.ResourceInUseException): + con.description # noqa: B018 + time.sleep(0.0005) + except Exception as error: + failures.append(repr(error)) def execute_once(): try: @@ -816,6 +903,7 @@ def execute_once(): stop.set() thread.join(10) con.close() + assert failures == [] assert refused == 0, (ran, refused) assert ran > 0 @@ -828,16 +916,21 @@ def test_a_submitting_call_beside_a_reader_is_refused_at_once(self): con.create_function("hold", lambda x: time.sleep(0.3) or x, [BIGINT], BIGINT) stop = threading.Event() waited = [] + failures = [] + + def until_stopped(call): + try: + while not stop.is_set(): + with contextlib.suppress(duckdb.ResourceInUseException): + call() + except Exception as error: + failures.append(repr(error)) def read_until_stopped(): - while not stop.is_set(): - with contextlib.suppress(duckdb.ResourceInUseException): - con.description # noqa: B018 + until_stopped(lambda: con.description) def run_held_statements(): - while not stop.is_set(): - with contextlib.suppress(duckdb.ResourceInUseException): - con.execute("SELECT hold(i) FROM range(1) t(i)") + until_stopped(lambda: con.execute("SELECT hold(i) FROM range(1) t(i)")) def execute_once(): began = time.monotonic() @@ -859,8 +952,10 @@ def execute_once(): for thread in threads: thread.join(10) con.close() + assert failures == [] assert waited == [] + @needs_pyarrow def test_other_calls_run_while_a_result_is_converted_to_arrow(self): # The Arrow fetch holds the connection only until the rows are taken and the schema is built; the conversion # into pyarrow that follows holds the mutex and no claim, so a relation runs during it and an execute waits @@ -881,16 +976,19 @@ def attempt(call): return "ran" def second(): - assert converting.wait(10) - deadline = time.monotonic() + 20 - while time.monotonic() < deadline and "done" not in marks: - began = time.monotonic() - kind = attempt(lambda: con.table("t")) - attempts.append((began, kind)) - # A relation admitted after the take was refused ran during the conversion; an execute then too - if kind == "ran" and any(k != "ran" for _, k in attempts): - outcome["execute during the conversion"] = attempt(lambda: con.execute("SELECT 1")) - break + try: + assert converting.wait(10) + deadline = time.monotonic() + 20 + while time.monotonic() < deadline and "done" not in marks: + began = time.monotonic() + kind = attempt(lambda: con.table("t")) + attempts.append((began, kind)) + # A relation admitted after the take was refused ran during the conversion; an execute then too + if kind == "ran" and any(k != "ran" for _, k in attempts): + outcome["execute during the conversion"] = attempt(lambda: con.execute("SELECT 1")) + break + except Exception as error: + outcome["second thread failed"] = repr(error) def fetch_or_none(): try: @@ -904,13 +1002,18 @@ def fetch_or_none(): con.execute("SELECT i, i::VARCHAR AS s FROM range(1000000) t(i)") converting.set() table = None - while table is None: + # Bounded: a fetch that lost its rows to a refusal returns None forever + give_up = time.monotonic() + 20 + for _ in range(1000): table = fetch_or_none() + if table is not None or time.monotonic() > give_up: + break marks["done"] = time.monotonic() finally: marks.setdefault("done", time.monotonic()) thread.join(30) con.close() + assert table is not None, (len(attempts), outcome) assert table.num_rows == 1000000 refused = [t for t, k in attempts if k != "ran"] assert refused, attempts[:5] From efe893aa163e425c83ef8a27af5233654f104fbe Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Wed, 11 Nov 2026 01:45:11 +0000 Subject: [PATCH 5/6] Supersede a result the parameter code registered before execute's claim, so the registration assertion states the rule --- src/pyconnection.cpp | 6 +++++- tests/fast/api/test_execute_contract.py | 12 +++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index 40623c93..d4e20dd4 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -889,6 +889,9 @@ std::shared_ptr DuckDBPyConnection::Execute(const nb::object } auto session = EngineSession(); auto named_values = TransformPreparedParameters(session.Context(), params); + // Converting the parameters runs their Python code, which may have run its own statement on this connection + // from this thread; that result gives way to this call's, as any earlier result does + SupersedeResult(false); DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; unique_lock lock(py_connection_lock); @@ -918,7 +921,8 @@ bool DuckDBPyConnection::RegisterResult(shared_ptr res) { if (!res) { return false; } - // The result the call superseded is gone already, so nothing that needs the GIL is destroyed here + // The call superseded the earlier result, and the one its parameter code may have registered, before the claim, + // so nothing that needs the GIL is destroyed here D_ASSERT(!con.HasResult()); con.SetResult(make_uniq(std::make_shared(std::move(res)))); return true; diff --git a/tests/fast/api/test_execute_contract.py b/tests/fast/api/test_execute_contract.py index f0fe1e07..5fcdb80a 100644 --- a/tests/fast/api/test_execute_contract.py +++ b/tests/fast/api/test_execute_contract.py @@ -801,7 +801,8 @@ def second(): def test_a_nested_execute_from_the_submitting_thread_still_enters(self): # The thread whose submitting call is in flight may submit again from inside it, as the recursive mutex - # admits it: a parameter object that runs a statement reads its row, and the outer call returns its own + # admits it: a parameter object that runs a statement reads its row, and the outer call's result then + # supersedes the nested one, so the outer call returns its own row and the connection holds only that result con = duckdb.connect() con.execute("SET threads = 4") inner = [] @@ -811,15 +812,20 @@ def __len__(self) -> int: return 1 def __iter__(self) -> Iterator[int]: - inner.append(con.execute("SELECT 99").fetchall()) + inner.append(con.execute("SELECT 99 AS nested").fetchall()) return iter([1]) try: - rows = con.execute("SELECT ?::INTEGER", Params()).fetchall() + con.execute("SELECT ?::INTEGER AS outer", Params()) + description = [column[0] for column in con.description] + rows = con.fetchall() + left = con.fetchall() finally: con.close() assert inner == [[(99,)]] + assert description == ["outer"] assert rows == [(1,)] + assert left == [] def test_a_reader_is_refused_while_a_statement_runs(self): # A call that waits for the mutex when a reader holds it is still refused outright while a statement holds From b54a0f1f22e612da6c40917bcb11498cacc14237 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Wed, 11 Nov 2026 02:48:22 +0000 Subject: [PATCH 6/6] Pin the unread nested result case and retry the Arrow fetch test by time --- .../pyconnection/pyconnection.hpp | 2 +- src/include/duckdb_python/pyresult.hpp | 2 +- src/pyconnection.cpp | 7 ++--- tests/fast/api/test_appender.py | 4 +-- tests/fast/api/test_execute_contract.py | 29 ++++++++++++++++--- 5 files changed, 32 insertions(+), 12 deletions(-) diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 01444999..829ce3c0 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -417,7 +417,7 @@ struct DuckDBPyConnection : public std::enable_shared_from_this Execute(const nb::object &query, nb::object params = nb::list()); - //! Makes the result the connection's DB-API result, under the caller's claim; false when there is none + //! Makes the result the connection's DB-API result, with the GIL released; false when there is none bool RegisterResult(shared_ptr res); std::shared_ptr ExecuteFromString(const string &query); diff --git a/src/include/duckdb_python/pyresult.hpp b/src/include/duckdb_python/pyresult.hpp index bb1e1da6..8ae1114e 100644 --- a/src/include/duckdb_python/pyresult.hpp +++ b/src/include/duckdb_python/pyresult.hpp @@ -129,7 +129,7 @@ struct DuckDBPyResult { duckdb::pyarrow::Table ConvertToArrowTable(const ArrowSchema &schema); duckdb::pyarrow::RecordBatchReader ImportRecordBatchReader(ArrowSchema *prebuilt, idx_t rows_per_batch); //! The stream's private data owns the result, and its callbacks run without the GIL. - //! prebuilt, when given, becomes the stream's schema; otherwise the result's own is copied or built + //! prebuilt, when given, becomes the stream's schema; otherwise the result's own is copied ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch, ArrowSchema *prebuilt); private: diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index d4e20dd4..9ccd6562 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -1,5 +1,7 @@ #include "duckdb_python/pyconnection/pyconnection.hpp" +#include "duckdb/common/identifier.hpp" + #include "duckdb/catalog/catalog.hpp" #include "duckdb/common/arrow/arrow.hpp" #include "duckdb/common/types.hpp" @@ -1571,10 +1573,7 @@ nb::object DuckDBPyConnection::ResultAsArrow(idx_t batch_size, bool reader) { if (!rows) { return nb::none(); } - vector schema_names; - for (auto &name : names) { - schema_names.push_back(name.GetIdentifierName()); - } + auto schema_names = IdentifiersToStrings(names); { DUCKDB_PY_ASSERT_GIL_HELD(); nb::gil_scoped_release release; diff --git a/tests/fast/api/test_appender.py b/tests/fast/api/test_appender.py index 3005f633..8e2d2ac6 100644 --- a/tests/fast/api/test_appender.py +++ b/tests/fast/api/test_appender.py @@ -1796,8 +1796,8 @@ def flush_until_stopped(): attempts = 0 try: while not stop.is_set(): - app.append_row(flushed, "a") try: + app.append_row(flushed, "a") app.flush() flushed += 1 except duckdb.ResourceInUseException as error: @@ -1875,8 +1875,8 @@ def flush_two_hundred_times(): go.wait(10) try: for attempt in range(200): - app.append_row(flushed, "a") try: + app.append_row(flushed, "a") app.flush() flushed += 1 except duckdb.ResourceInUseException as error: diff --git a/tests/fast/api/test_execute_contract.py b/tests/fast/api/test_execute_contract.py index 5fcdb80a..1ca5929c 100644 --- a/tests/fast/api/test_execute_contract.py +++ b/tests/fast/api/test_execute_contract.py @@ -726,18 +726,24 @@ def bind_until_stopped(): con.execute("SELECT * FROM range(5) t(i)") except duckdb.ResourceInUseException: continue - table = None - for attempt in range(1000): + # Retried by time: the binding thread can refuse the fetch many times in a row under load + table, admitted, attempt = None, False, 0 + give_up = time.monotonic() + 5 + while time.monotonic() < give_up: try: table = con.to_arrow_table() + admitted = True break except duckdb.ResourceInUseException as error: + attempt += 1 if {BUSY!r} not in str(error): counts["refused: " + str(error)] += 1 fetches += 1 if attempt: retried += 1 - if table is None: + if not admitted: + counts["fetch never admitted in 5 s"] += 1 + elif table is None: counts["no table" + (" after a refusal" if attempt else "")] += 1 elif table.num_rows != 5: counts[f"rows {{table.num_rows}}" + (" after a refusal" if attempt else "")] += 1 @@ -802,7 +808,8 @@ def second(): def test_a_nested_execute_from_the_submitting_thread_still_enters(self): # The thread whose submitting call is in flight may submit again from inside it, as the recursive mutex # admits it: a parameter object that runs a statement reads its row, and the outer call's result then - # supersedes the nested one, so the outer call returns its own row and the connection holds only that result + # supersedes the nested one, read out or left open, so the outer call returns its own row and the connection + # holds only that result con = duckdb.connect() con.execute("SET threads = 4") inner = [] @@ -815,17 +822,31 @@ def __iter__(self) -> Iterator[int]: inner.append(con.execute("SELECT 99 AS nested").fetchall()) return iter([1]) + class UnreadParams: + def __len__(self) -> int: + return 1 + + def __iter__(self) -> Iterator[int]: + # Left open on purpose: the outer call supersedes it, where it used to be refused for it + con.execute("SELECT 99 AS nested") + return iter([2]) + try: con.execute("SELECT ?::INTEGER AS outer", Params()) description = [column[0] for column in con.description] rows = con.fetchall() left = con.fetchall() + con.execute("SELECT ?::INTEGER AS outer", UnreadParams()) + unread_description = [column[0] for column in con.description] + unread_rows = con.fetchall() finally: con.close() assert inner == [[(99,)]] assert description == ["outer"] assert rows == [(1,)] assert left == [] + assert unread_description == ["outer"] + assert unread_rows == [(2,)] def test_a_reader_is_refused_while_a_statement_runs(self): # A call that waits for the mutex when a reader holds it is still refused outright while a statement holds