diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 638b0a4b..e5bcf7a3 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -321,7 +321,7 @@ struct DuckDBPyConnection : public std::enable_shared_from_this GetDescription(); - int GetRowcount(); + int64_t GetRowcount(); // these should be functions on the result but well Optional FetchOne(); diff --git a/src/include/duckdb_python/pyrelation.hpp b/src/include/duckdb_python/pyrelation.hpp index f71a6327..b1611026 100644 --- a/src/include/duckdb_python/pyrelation.hpp +++ b/src/include/duckdb_python/pyrelation.hpp @@ -30,6 +30,15 @@ struct DuckDBPyRelation { nb::list Description(); + //! Cached at construction time (see the DuckDBPyResult constructor overload) so that it survives + //! the Fetch*() methods below, several of which null out `result` once fully consumed. Note this + //! is only ever non-default for relations built directly from a DuckDBPyResult (the connection's + //! execute()/executemany() path) - relations built from a lazy Relation (rel != nullptr, executed + //! via ExecuteOrThrow) are always SELECT-shaped and so never produce a CHANGED_ROWS value anyway. + int64_t GetRowcount() const { + return row_changes; + } + void Close(); std::unique_ptr GetAttribute(const string &name); @@ -300,6 +309,8 @@ struct DuckDBPyRelation { vector names; std::shared_ptr result; std::string rendered_result; + //! Cached row-changed count - see GetRowcount(). + int64_t row_changes = -1; }; } // namespace duckdb diff --git a/src/include/duckdb_python/pyresult.hpp b/src/include/duckdb_python/pyresult.hpp index 865f955f..16a3d69a 100644 --- a/src/include/duckdb_python/pyresult.hpp +++ b/src/include/duckdb_python/pyresult.hpp @@ -59,6 +59,12 @@ struct DuckDBPyResult { ClientProperties GetClientProperties(); + //! Number of rows changed by the last CHANGED_ROWS-returning statement (INSERT/UPDATE/DELETE/...). + //! Returns -1 when not applicable/unknown, as permitted by the DB-API 2.0 spec for 'rowcount'. + int64_t GetRowcount() const { + return row_changes; + } + private: void FillNumpy(nb::dict &res, idx_t col_idx, NumpyResultConversion &conversion, const char *name); @@ -82,6 +88,12 @@ struct DuckDBPyResult { duckdb::pyarrow::Table MaterializedResultToArrowTable(const ArrowSchema &arrow_schema, idx_t rows_per_batch); ArrowArrayStream FetchArrowArrayStream(idx_t rows_per_batch); + //! Computes the CHANGED_ROWS value (if any) up front, before any Fetch call has had a chance to + //! consume it, so that later fetch calls do not affect what GetRowcount() reports. Materializes a + //! streaming result if necessary; this is cheap since CHANGED_ROWS results are always exactly one + //! already-computed row. + int64_t ComputeRowChanges(); + private: idx_t chunk_offset = 0; @@ -92,6 +104,8 @@ struct DuckDBPyResult { // Holds the categorical type of Categorical/ENUM types unordered_map categories_type; bool result_closed = false; + //! Cached by ComputeRowChanges() at construction time - see GetRowcount(). + int64_t row_changes = -1; }; } // namespace duckdb diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index ebdf25fa..489379cd 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -1920,8 +1920,13 @@ Optional DuckDBPyConnection::GetDescription() { return result.Description(); } -int DuckDBPyConnection::GetRowcount() { - return -1; +int64_t DuckDBPyConnection::GetRowcount() { + ConnectionLockGuard conn_lock(*this); + if (!con.HasResult()) { + return -1; + } + auto &result = con.GetResult(); + return result.GetRowcount(); } void DuckDBPyConnection::Close() { diff --git a/src/pyrelation.cpp b/src/pyrelation.cpp index 2d1ac47b..ff894dfb 100644 --- a/src/pyrelation.cpp +++ b/src/pyrelation.cpp @@ -71,6 +71,7 @@ DuckDBPyRelation::DuckDBPyRelation(std::shared_ptr result_p) this->executed = true; this->types = result->GetTypes(); this->names = result->GetNames(); + this->row_changes = result->GetRowcount(); } std::unique_ptr DuckDBPyRelation::ProjectFromExpression(const string &expression) { diff --git a/src/pyresult.cpp b/src/pyresult.cpp index 7f0d0c9a..db3d6c0b 100644 --- a/src/pyresult.cpp +++ b/src/pyresult.cpp @@ -33,6 +33,9 @@ DuckDBPyResult::DuckDBPyResult(unique_ptr result_p) : result(std::m if (!result) { throw InternalException("PyResult created without a result object"); } + // Must happen before any Fetch call (from this object or a caller holding a reference to the + // same underlying QueryResult) has a chance to consume the single row we're reading here. + row_changes = ComputeRowChanges(); } DuckDBPyResult::~DuckDBPyResult() { @@ -68,6 +71,43 @@ const vector &DuckDBPyResult::GetTypes() { return result->types; } +int64_t DuckDBPyResult::ComputeRowChanges() { + if (!result || result->HasError()) { + return -1; + } + if (result->properties.return_type != StatementReturnType::CHANGED_ROWS) { + // The row count of a SELECT (or a statement that returns nothing) is not known without fully + // consuming the result - report -1, as permitted by the DB-API 2.0 spec for 'rowcount'. + return -1; + } + if (result->type == QueryResultType::STREAM_RESULT) { + // CHANGED_ROWS statements always produce a single already-computed row, so materializing here + // does not trigger any additional query execution - it just changes the in-memory representation. + // Still release the GIL around it though: Materialize() drives the same native Fetch() machinery + // as any other result consumption, and other fetch paths (e.g. Fetchone()) release the GIL around it. + unique_ptr materialized; + { + D_ASSERT(duckdb::PyUtil::GilCheck()); + nb::gil_scoped_release release; + auto &stream_result = result->Cast(); + materialized = stream_result.Materialize(); + } + if (!materialized || materialized->HasError()) { + return -1; + } + result = std::move(materialized); + } + if (result->type != QueryResultType::MATERIALIZED_RESULT) { + // e.g. an Arrow result - can't peek at the value without disturbing it. + return -1; + } + auto &materialized_result = result->Cast(); + if (materialized_result.RowCount() != 1 || materialized_result.ColumnCount() != 1) { + return -1; + } + return materialized_result.GetValue(0, 0).GetValue(); +} + unique_ptr DuckDBPyResult::FetchChunk() { if (!result) { throw InternalException("FetchChunk called without a result object"); diff --git a/tests/fast/api/test_dbapi10.py b/tests/fast/api/test_dbapi10.py index 6d60b27c..a3907bc5 100644 --- a/tests/fast/api/test_dbapi10.py +++ b/tests/fast/api/test_dbapi10.py @@ -56,3 +56,104 @@ def test_none_description(self, duckdb_empty_cursor): class TestCursorRowcount: def test_rowcount(self, duckdb_cursor): assert duckdb_cursor.rowcount == -1 + + def test_rowcount_no_query_yet(self, duckdb_cursor): + assert duckdb_cursor.rowcount == -1 + + def test_rowcount_insert(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3)") + assert duckdb_cursor.rowcount == 3 + + def test_rowcount_insert_select(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t SELECT * FROM range(100)") + assert duckdb_cursor.rowcount == 100 + + def test_rowcount_update(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t AS SELECT * FROM range(10) t(i)") + duckdb_cursor.execute("UPDATE t SET i = i + 1 WHERE i < 4") + assert duckdb_cursor.rowcount == 4 + + def test_rowcount_delete(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t AS SELECT * FROM range(10) t(i)") + duckdb_cursor.execute("DELETE FROM t WHERE i < 3") + assert duckdb_cursor.rowcount == 3 + + def test_rowcount_create_table_as(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t AS SELECT * FROM range(42) t(i)") + assert duckdb_cursor.rowcount == 42 + + def test_rowcount_select_is_unknown(self, duckdb_cursor): + # Matches DB-API 2.0: rowcount is -1 when it can't be determined without consuming the result. + duckdb_cursor.execute("SELECT * FROM range(10)") + assert duckdb_cursor.rowcount == -1 + duckdb_cursor.fetchall() + assert duckdb_cursor.rowcount == -1 + + def test_rowcount_reset_by_next_execute(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2)") + assert duckdb_cursor.rowcount == 2 + duckdb_cursor.execute("SELECT * FROM t") + assert duckdb_cursor.rowcount == -1 + + def test_rowcount_does_not_disturb_fetch(self, duckdb_cursor): + # Accessing rowcount must not consume the single-row result that fetchone() also reads. + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3)") + assert duckdb_cursor.rowcount == 3 + assert duckdb_cursor.fetchone() == (3,) + + def test_rowcount_ddl_is_unknown(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + assert duckdb_cursor.rowcount == -1 + + def test_rowcount_after_fetchall(self, duckdb_cursor): + # Regression test: rowcount must survive full consumption of the result via fetchall(), + # not just a bare execute() with no fetch at all. + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)") + assert duckdb_cursor.fetchall() == [(4,)] + assert duckdb_cursor.rowcount == 4 + + def test_rowcount_after_fetchone(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)") + assert duckdb_cursor.fetchone() == (4,) + assert duckdb_cursor.rowcount == 4 + + def test_rowcount_after_fetchmany(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)") + assert duckdb_cursor.fetchmany(1) == [(4,)] + assert duckdb_cursor.rowcount == 4 + + def test_rowcount_after_fetchdf(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)") + duckdb_cursor.fetchdf() + assert duckdb_cursor.rowcount == 4 + + def test_rowcount_after_fetchnumpy(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)") + duckdb_cursor.fetchnumpy() + assert duckdb_cursor.rowcount == 4 + + def test_rowcount_after_to_arrow_table(self, duckdb_cursor): + pytest.importorskip("pyarrow") + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.execute("INSERT INTO t VALUES (1), (2), (3), (4)") + duckdb_cursor.to_arrow_table() + assert duckdb_cursor.rowcount == 4 + + def test_rowcount_executemany_reflects_last_statement_only(self, duckdb_cursor): + # Documents current (limited) behavior: executemany() only keeps the QueryResult of the last + # parameter set it executes, so rowcount reflects that one statement rather than the total + # number of rows affected across all parameter sets. Not the DB-API-idiomatic answer, but + # pinned here so a future change to this behavior is a deliberate, visible decision. + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + duckdb_cursor.executemany("INSERT INTO t VALUES (?)", [(1,), (2,), (3,)]) + assert duckdb_cursor.table("t").fetchall() == [(1,), (2,), (3,)] + assert duckdb_cursor.rowcount == 1