diff --git a/_duckdb-stubs/__init__.pyi b/_duckdb-stubs/__init__.pyi index 8770483f..4d5ff6d4 100644 --- a/_duckdb-stubs/__init__.pyi +++ b/_duckdb-stubs/__init__.pyi @@ -46,6 +46,7 @@ if typing.TYPE_CHECKING: from duckdb import sqltypes, func __all__: lst[str] = [ + "Appender", "BinderException", "CSVLineTerminator", "CaseExpression", @@ -201,6 +202,7 @@ class DuckDBPyConnection: def __enter__(self) -> Self: ... def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: ... def append(self, table_name: str, df: pandas.DataFrame, *, by_name: bool = False) -> DuckDBPyConnection: ... + def appender(self, table: str, schema: str | None = None, catalog: str | None = None) -> Appender: ... def array_type(self, type: IntoPyType, size: typing.SupportsInt) -> sqltypes.DuckDBPyType: ... def arrow(self, rows_per_batch: typing.SupportsInt = 1000000) -> pyarrow.lib.RecordBatchReader: """Alias of to_arrow_reader(). We recommend using to_arrow_reader() instead.""" @@ -841,6 +843,16 @@ class ProgrammingError(DatabaseError): ... class SequenceException(DatabaseError): ... class SerializationException(OperationalError): ... +class Appender: + def __enter__(self) -> Self: ... + def __exit__(self, exc_type: object, exc: object, traceback: object) -> None: ... + def append(self, *args: object) -> None: ... + def append_chunk(self, rows: Sequence[Sequence[object]]) -> None: ... + def close(self) -> None: ... + def flush(self) -> None: ... + @property + def column_count(self) -> int: ... + class Statement: @property def expected_result_type(self) -> lst[StatementType]: ... diff --git a/duckdb/__init__.py b/duckdb/__init__.py index d17c530f..fb8f1048 100644 --- a/duckdb/__init__.py +++ b/duckdb/__init__.py @@ -10,6 +10,7 @@ """ from _duckdb import ( + Appender, BinderException, CaseExpression, CatalogException, @@ -207,6 +208,7 @@ ) __all__: list[str] = [ + "Appender", "BinaryValue", "BinderException", "BitValue", diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 59418aa5..c0018fa0 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -18,6 +18,7 @@ add_library( importer.cpp map.cpp path_like.cpp + pyappender.cpp pyconnection.cpp pyexpression.cpp pyfilesystem.cpp diff --git a/src/duckdb_python.cpp b/src/duckdb_python.cpp index af0f7abe..24e2856a 100644 --- a/src/duckdb_python.cpp +++ b/src/duckdb_python.cpp @@ -7,6 +7,7 @@ #include "duckdb_python/python_objects.hpp" #include "duckdb_python/pyconnection/pyconnection.hpp" #include "duckdb_python/pystatement.hpp" +#include "duckdb_python/pyappender.hpp" #include "duckdb_python/pyrelation.hpp" #include "duckdb_python/expression/pyexpression.hpp" #include "duckdb_python/exceptions.hpp" @@ -1108,6 +1109,7 @@ NB_MODULE(DUCKDB_PYTHON_LIB_NAME, m) { // NOLINT DuckDBPyExpression::Initialize(m); DuckDBPyStatement::Initialize(m); DuckDBPyRelation::Initialize(m); + DuckDBPyAppender::Initialize(m); DuckDBPyConnection::Initialize(m); PythonObject::Initialize(); diff --git a/src/include/duckdb_python/pyappender.hpp b/src/include/duckdb_python/pyappender.hpp new file mode 100644 index 00000000..fcce4f88 --- /dev/null +++ b/src/include/duckdb_python/pyappender.hpp @@ -0,0 +1,42 @@ +//===----------------------------------------------------------------------===// +// DuckDB +// +// duckdb_python/pyappender.hpp +// +// +//===----------------------------------------------------------------------===// + +#pragma once + +#include "duckdb_python/nb/casters.hpp" +#include "duckdb.hpp" +#include "duckdb_python/pyconnection/pyconnection.hpp" + +#include + +namespace duckdb { + +struct DuckDBPyAppender : std::enable_shared_from_this { +public: + DuckDBPyAppender(std::shared_ptr connection, unique_ptr appender); + ~DuckDBPyAppender(); + + static void Initialize(nb::handle &m); + + void AppendRow(const nb::args &args); + void AppendChunk(const nb::object &rows); + void Flush(); + void Close(); + std::shared_ptr Enter(); + void Exit(const nb::object &exc_type, const nb::object &exc, const nb::object &traceback); + idx_t ColumnCount(); + +private: + void CheckOpen() const; + ClientContext &Context(); + + std::shared_ptr connection; + unique_ptr appender; +}; + +} // namespace duckdb diff --git a/src/include/duckdb_python/pyconnection/pyconnection.hpp b/src/include/duckdb_python/pyconnection/pyconnection.hpp index 638b0a4b..521ff166 100644 --- a/src/include/duckdb_python/pyconnection/pyconnection.hpp +++ b/src/include/duckdb_python/pyconnection/pyconnection.hpp @@ -24,12 +24,15 @@ #include "duckdb_python/nb/conversions/python_udf_type_enum.hpp" #include "duckdb/common/shared_ptr.hpp" +#include + namespace duckdb { struct BoundParameterData; enum class PythonEnvironmentType { NORMAL, INTERACTIVE, JUPYTER }; struct DuckDBPyRelation; +struct DuckDBPyAppender; class RegisteredArrow : public RegisteredObject { @@ -269,6 +272,9 @@ struct DuckDBPyConnection : public std::enable_shared_from_this Append(const string &name, const PandasDataFrame &value, bool by_name); + std::shared_ptr CreateAppender(const string &table, std::optional schema = std::nullopt, + std::optional catalog = std::nullopt); + std::shared_ptr RegisterPythonObject(const string &name, const nb::object &python_object); void InstallExtension(const string &extension, bool force_install = false, diff --git a/src/pyappender.cpp b/src/pyappender.cpp new file mode 100644 index 00000000..650622b3 --- /dev/null +++ b/src/pyappender.cpp @@ -0,0 +1,161 @@ +#include "duckdb_python/pyappender.hpp" +#include "duckdb_python/python_conversion.hpp" + +namespace duckdb { + +DuckDBPyAppender::DuckDBPyAppender(std::shared_ptr connection_p, + unique_ptr appender_p) + : connection(std::move(connection_p)), appender(std::move(appender_p)) { +} + +DuckDBPyAppender::~DuckDBPyAppender() { + if (!appender) { + return; + } + try { + appender->Close(); + } catch (...) { // NOLINT + } + appender.reset(); +} + +void DuckDBPyAppender::Initialize(nb::handle &m) { + auto appender_module = nb::class_(m, "Appender", nb::is_weak_referenceable()); + appender_module.def("append", &DuckDBPyAppender::AppendRow, "Append a row of values") + .def("append_chunk", &DuckDBPyAppender::AppendChunk, "Append a sequence of rows as a data chunk", + nb::arg("rows")) + .def("flush", &DuckDBPyAppender::Flush, "Flush the appender to the table") + .def("close", &DuckDBPyAppender::Close, "Flush the appender and close it") + .def("__enter__", &DuckDBPyAppender::Enter) + .def("__exit__", &DuckDBPyAppender::Exit, nb::arg("exc_type").none(), nb::arg("exc").none(), + nb::arg("traceback").none()) + .def_prop_ro("column_count", &DuckDBPyAppender::ColumnCount, "Number of columns in the appender"); +} + +void DuckDBPyAppender::CheckOpen() const { + if (!appender) { + throw InvalidInputException("This appender has been closed"); + } +} + +ClientContext &DuckDBPyAppender::Context() { + return *connection->con.GetConnection().context; +} + +void DuckDBPyAppender::AppendRow(const nb::args &args) { + DuckDBPyConnection::ConnectionLockGuard conn_lock(*connection); + CheckOpen(); + auto &types = appender->GetActiveTypes(); + if (args.size() != types.size()) { + throw InvalidInputException("appender.append expected %d values, got %d", types.size(), args.size()); + } + vector values; + values.reserve(types.size()); + idx_t i = 0; + for (auto value : args) { + values.push_back(TransformPythonValue(Context(), value, types[i])); + i++; + } + appender->BeginRow(); + for (auto &value : values) { + appender->Append(std::move(value)); + } + appender->EndRow(); +} + +void DuckDBPyAppender::AppendChunk(const nb::object &rows) { + DuckDBPyConnection::ConnectionLockGuard conn_lock(*connection); + CheckOpen(); + if (!duckdb::PyUtil::IsListLike(rows)) { + throw InvalidInputException("append_chunk expects a sequence of rows"); + } + nb::list list(rows); + auto &types = appender->GetActiveTypes(); + auto &context = Context(); + idx_t offset = 0; + const idx_t total = list.size(); + while (offset < total) { + const idx_t n = MinValue(total - offset, STANDARD_VECTOR_SIZE); + DataChunk chunk; + chunk.Initialize(context, types); + chunk.SetChildCardinality(n); + for (idx_t r = 0; r < n; r++) { + auto row_obj = list[offset + r]; + if (!duckdb::PyUtil::IsListLike(row_obj)) { + throw InvalidInputException("each append_chunk row must be a sequence"); + } + nb::list row(row_obj); + if (row.size() != types.size()) { + throw InvalidInputException("append_chunk row expected %d values, got %d", types.size(), row.size()); + } + idx_t c = 0; + for (auto value : row) { + TransformPythonObject(&context, value, chunk.data[c], r); + c++; + } + } + { + nb::gil_scoped_release release; + appender->AppendDataChunk(chunk); + } + offset += n; + } +} + +void DuckDBPyAppender::Flush() { + DuckDBPyConnection::ConnectionLockGuard conn_lock(*connection); + CheckOpen(); + nb::gil_scoped_release release; + appender->Flush(); +} + +void DuckDBPyAppender::Close() { + DuckDBPyConnection::ConnectionLockGuard conn_lock(*connection); + if (!appender) { + return; + } + { + nb::gil_scoped_release release; + appender->Close(); + } + appender.reset(); +} + +std::shared_ptr DuckDBPyAppender::Enter() { + DuckDBPyConnection::ConnectionLockGuard conn_lock(*connection); + CheckOpen(); + return shared_from_this(); +} + +void DuckDBPyAppender::Exit(const nb::object &, const nb::object &, const nb::object &) { + Close(); +} + +idx_t DuckDBPyAppender::ColumnCount() { + DuckDBPyConnection::ConnectionLockGuard conn_lock(*connection); + CheckOpen(); + return appender->GetActiveTypes().size(); +} + +std::shared_ptr DuckDBPyConnection::CreateAppender(const string &table, std::optional schema, + std::optional catalog) { + if (catalog.has_value() && !schema.has_value()) { + throw InvalidInputException("catalog requires schema"); + } + DuckDBPyConnection::ConnectionLockGuard conn_lock(*this); + auto &con = this->con.GetConnection(); + unique_ptr appender; + { + nb::gil_scoped_release release; + if (catalog.has_value()) { + appender = make_uniq(con, Identifier(*catalog), Identifier(*schema), Identifier(table)); + } else if (schema.has_value()) { + appender = make_uniq(con, Identifier(*schema), Identifier(table)); + } else { + appender = make_uniq(con, Identifier(table)); + } + } + return std::make_shared(shared_from_this(), std::move(appender)); +} + +} // namespace duckdb diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index ebdf25fa..0659ce7f 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -1,4 +1,5 @@ #include "duckdb_python/pyconnection/pyconnection.hpp" +#include "duckdb_python/pyappender.hpp" #include "duckdb/catalog/catalog.hpp" #include "duckdb/common/arrow/arrow.hpp" @@ -502,6 +503,9 @@ void DuckDBPyConnection::Initialize(nb::handle &m) { connection_module.def("__del__", &DuckDBPyConnection::Close); InitializeConnectionMethods(connection_module); + connection_module.def("appender", &DuckDBPyConnection::CreateAppender, + "Create an appender for fast row-wise inserts into a table", nb::arg("table"), + nb::arg("schema") = nb::none(), nb::arg("catalog") = nb::none()); connection_module.def_prop_ro("description", &DuckDBPyConnection::GetDescription, "Get result set attributes, mainly column names"); connection_module.def_prop_ro("rowcount", &DuckDBPyConnection::GetRowcount, "Get result set row count"); diff --git a/tests/fast/test_appender.py b/tests/fast/test_appender.py new file mode 100644 index 00000000..a7249ff2 --- /dev/null +++ b/tests/fast/test_appender.py @@ -0,0 +1,82 @@ +import datetime + +import pytest + +import duckdb + + +class TestAppender: + def test_create_append_flush_close(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE people (id INTEGER, name VARCHAR)") + appender = duckdb_cursor.appender("people") + appender.append(1, "Mark") + appender.append(2, "Hannes") + appender.flush() + assert duckdb_cursor.execute("SELECT * FROM people ORDER BY id").fetchall() == [(1, "Mark"), (2, "Hannes")] + appender.close() + + def test_close_flushes(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + appender = duckdb_cursor.appender("t") + appender.append(42) + appender.close() + assert duckdb_cursor.execute("SELECT i FROM t").fetchall() == [(42,)] + + def test_append_after_close_raises(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + appender = duckdb_cursor.appender("t") + appender.close() + with pytest.raises(duckdb.Error): + appender.append(1) + + def test_schema_and_catalog(self, duckdb_cursor): + duckdb_cursor.execute("CREATE SCHEMA s") + duckdb_cursor.execute("CREATE TABLE s.t (i INTEGER, b BOOLEAN, d DOUBLE)") + appender = duckdb_cursor.appender("t", schema="s") + appender.append(7, True, 1.5) + appender.append(None, False, None) + appender.close() + assert duckdb_cursor.execute("SELECT * FROM s.t ORDER BY i NULLS LAST").fetchall() == [ + (7, True, 1.5), + (None, False, None), + ] + + def test_scalar_types(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (b BOOLEAN, i INTEGER, l BIGINT, f FLOAT, d DOUBLE, s VARCHAR, dt DATE)") + appender = duckdb_cursor.appender("t") + appender.append(True, 1, 2, 1.25, 2.5, "x", datetime.date(2026, 8, 16)) + appender.close() + row = duckdb_cursor.execute("SELECT * FROM t").fetchone() + assert row[0] is True + assert row[1] == 1 + assert row[2] == 2 + assert row[5] == "x" + assert row[6] == datetime.date(2026, 8, 16) + + def test_wrong_column_count(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER, j INTEGER)") + appender = duckdb_cursor.appender("t") + with pytest.raises(duckdb.Error): + appender.append(1) + appender.close() + + def test_conversion_error_does_not_drop_prior_rows(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER, j INTEGER)") + appender = duckdb_cursor.appender("t") + appender.append(1, 2) + with pytest.raises(duckdb.Error): + appender.append(3, "nope") + appender.append(4, 5) + appender.close() + assert duckdb_cursor.execute("SELECT * FROM t ORDER BY i").fetchall() == [(1, 2), (4, 5)] + + def test_missing_table(self, duckdb_cursor): + with pytest.raises(duckdb.Error): + duckdb_cursor.appender("no_such_table") + + def test_context_manager(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER)") + with duckdb_cursor.appender("t") as appender: + appender.append(1) + appender.append(2) + assert duckdb_cursor.execute("SELECT * FROM t ORDER BY i").fetchall() == [(1,), (2,)] diff --git a/tests/fast/test_appender_chunk.py b/tests/fast/test_appender_chunk.py new file mode 100644 index 00000000..187a0d02 --- /dev/null +++ b/tests/fast/test_appender_chunk.py @@ -0,0 +1,37 @@ +import duckdb + + +class TestAppenderChunk: + def test_append_chunk(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER, s VARCHAR)") + appender = duckdb_cursor.appender("t") + appender.append_chunk([[1, "a"], [2, "b"], [3, "c"]]) + appender.close() + assert duckdb_cursor.execute("SELECT * FROM t ORDER BY i").fetchall() == [(1, "a"), (2, "b"), (3, "c")] + + def test_append_chunk_wrong_width(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER, j INTEGER)") + appender = duckdb_cursor.appender("t") + try: + appender.append_chunk([[1]]) + raise AssertionError("expected error") + except duckdb.Error: + pass + appender.close() + + +class TestAppenderNested: + def test_list(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (i INTEGER, xs INTEGER[])") + appender = duckdb_cursor.appender("t") + appender.append(1, [10, 20, 30]) + appender.append(2, []) + appender.close() + assert duckdb_cursor.execute("SELECT * FROM t ORDER BY i").fetchall() == [(1, [10, 20, 30]), (2, [])] + + def test_struct(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (x STRUCT(a INTEGER, b VARCHAR))") + appender = duckdb_cursor.appender("t") + appender.append({"a": 1, "b": "x"}) + appender.close() + assert duckdb_cursor.execute("SELECT x FROM t").fetchall() == [({"a": 1, "b": "x"},)]