diff --git a/_duckdb-stubs/__init__.pyi b/_duckdb-stubs/__init__.pyi index 81d69be7..f8a15607 100644 --- a/_duckdb-stubs/__init__.pyi +++ b/_duckdb-stubs/__init__.pyi @@ -547,7 +547,7 @@ class DuckDBPyRelation: def count( self, expression: str, groups: str = "", window_spec: str = "", projected_columns: str = "" ) -> DuckDBPyRelation: ... - def create(self, table_name: str) -> None: ... + def create(self, table_name: str, replace: bool = False) -> None: ... def create_view(self, view_name: str, replace: bool = True) -> DuckDBPyRelation: ... def cross(self, other_rel: DuckDBPyRelation) -> DuckDBPyRelation: ... def cume_dist(self, window_spec: str, projected_columns: str = "") -> DuckDBPyRelation: ... diff --git a/src/duckdb_py/include/duckdb_python/pyrelation.hpp b/src/duckdb_py/include/duckdb_python/pyrelation.hpp index 50f39b5f..d17e0a62 100644 --- a/src/duckdb_py/include/duckdb_python/pyrelation.hpp +++ b/src/duckdb_py/include/duckdb_python/pyrelation.hpp @@ -239,7 +239,7 @@ struct DuckDBPyRelation { void Insert(const py::object ¶ms = py::list()) const; void Update(const py::object &set, const py::object &where = py::none()); - void Create(const string &table); + void Create(const string &table, bool replace = false); py::str Type(); py::list Columns(); diff --git a/src/duckdb_py/pyrelation.cpp b/src/duckdb_py/pyrelation.cpp index 8e223cdb..b6a90cac 100644 --- a/src/duckdb_py/pyrelation.cpp +++ b/src/duckdb_py/pyrelation.cpp @@ -1,5 +1,6 @@ #include "duckdb_python/pybind11/pybind_wrapper.hpp" #include "duckdb_python/pyrelation.hpp" +#include "duckdb/common/enums/on_create_conflict.hpp" #include "duckdb_python/pyconnection/pyconnection.hpp" #include "duckdb_python/pytype.hpp" #include "duckdb_python/pyresult.hpp" @@ -1648,10 +1649,11 @@ void DuckDBPyRelation::Insert(const py::object ¶ms) const { rel->Insert(values); } -void DuckDBPyRelation::Create(const string &table) { +void DuckDBPyRelation::Create(const string &table, bool replace) { AssertRelation(); auto parsed_info = QualifiedName::Parse(table); - auto create = rel->CreateRel(parsed_info.schema, parsed_info.name, false); + auto on_conflict = replace ? OnCreateConflict::REPLACE_ON_CONFLICT : OnCreateConflict::ERROR_ON_CONFLICT; + auto create = rel->CreateRel(parsed_info.schema, parsed_info.name, false, on_conflict); PyExecuteRelation(create); } diff --git a/src/duckdb_py/pyrelation/initialize.cpp b/src/duckdb_py/pyrelation/initialize.cpp index 4393889a..f64bb4dd 100644 --- a/src/duckdb_py/pyrelation/initialize.cpp +++ b/src/duckdb_py/pyrelation/initialize.cpp @@ -330,7 +330,7 @@ void DuckDBPyRelation::Initialize(py::handle &m) { DefineMethod({"create", "to_table"}, relation_module, &DuckDBPyRelation::Create, "Creates a new table named table_name with the contents of the relation object", - py::arg("table_name")); + py::arg("table_name"), py::arg("replace") = false); DefineMethod({"create_view", "to_view"}, relation_module, &DuckDBPyRelation::CreateView, "Creates a view named view_name that refers to the relation object", py::arg("view_name"), diff --git a/tests/fast/test_relation.py b/tests/fast/test_relation.py index bc7039fa..706d9429 100644 --- a/tests/fast/test_relation.py +++ b/tests/fast/test_relation.py @@ -172,9 +172,11 @@ def test_except_operator(self): def test_create_operator(self): conn = duckdb.connect() + test_df = pd.DataFrame.from_dict({"i": [1, 2, 3, 4], "j": ["one", "two", "three", "four"]}) rel = conn.from_df(test_df) rel.create("test_df") + assert conn.query("select * from test_df").execute().fetchall() == [ (1, "one"), (2, "two"), @@ -182,6 +184,61 @@ def test_create_operator(self): (4, "four"), ] + def test_create_replace_false(self): + conn = duckdb.connect() + + test_df = pd.DataFrame.from_dict({"i": [1, 2, 3], "j": ["a", "b", "c"]}) + rel = conn.from_df(test_df) + rel.create("test_replace_tbl", replace=False) + + assert conn.query("select * from test_replace_tbl").execute().fetchall() == [ + (1, "a"), + (2, "b"), + (3, "c"), + ] + + def test_create_replace_true_different_schema(self): + conn = duckdb.connect() + + test_df = pd.DataFrame.from_dict({"i": [1, 2, 3], "j": ["a", "b", "c"]}) + rel = conn.from_df(test_df) + rel.create("test_replace_tbl", replace=True) + + assert conn.query("select * from test_replace_tbl").execute().fetchall() == [ + (1, "a"), + (2, "b"), + (3, "c"), + ] + + test_df2 = pd.DataFrame.from_dict({"a": [10, 20], "b": ["x", "y"], "c": [1.5, 2.5]}) + rel2 = conn.from_df(test_df2) + rel2.create("test_replace_tbl", replace=True) + + assert conn.query("select * from test_replace_tbl").execute().fetchall() == [ + (10, "x", 1.5), + (20, "y", 2.5), + ] + + def test_create_replace_false_error_on_existing(self): + conn = duckdb.connect() + + test_df = pd.DataFrame.from_dict({"i": [1], "j": ["a"]}) + rel = conn.from_df(test_df) + rel.create("test_replace_existing_tbl") + + with pytest.raises(duckdb.CatalogException): + rel.create("test_replace_existing_tbl", replace=False) + + def test_create_replace_true_error_on_non_existent(self): + conn = duckdb.connect() + + test_df = pd.DataFrame.from_dict({"i": [1], "j": ["a"]}) + rel = conn.from_df(test_df) + conn.execute("drop table if exists test_replace_tbl2") + rel.create("test_replace_tbl2", replace=True) + + assert conn.query("select * from test_replace_tbl2").execute().fetchall() == [(1, "a")] + def test_create_view_operator(self): conn = duckdb.connect() test_df = pd.DataFrame.from_dict({"i": [1, 2, 3, 4], "j": ["one", "two", "three", "four"]})