From 12db8cd8a80720693562858cb8535b930f04be80 Mon Sep 17 00:00:00 2001
From: Nick Crews <nicholas.b.crews@gmail.com>
Date: Fri, 23 Aug 2024 23:04:14 -0800
Subject: [PATCH] feat: support SchemaLike in Backend.create_table()

---
 ibis/backends/bigquery/__init__.py   |  2 ++
 ibis/backends/clickhouse/__init__.py |  7 +++----
 ibis/backends/duckdb/__init__.py     |  6 +++++-
 ibis/backends/exasol/__init__.py     |  2 ++
 ibis/backends/mssql/__init__.py      |  2 ++
 ibis/backends/tests/test_client.py   | 22 +++++-----------------
 6 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/ibis/backends/bigquery/__init__.py b/ibis/backends/bigquery/__init__.py
index a1bef8f57f2fa..6d097c4e3db14 100644
--- a/ibis/backends/bigquery/__init__.py
+++ b/ibis/backends/bigquery/__init__.py
@@ -972,6 +972,8 @@ def create_table(
         """
         if obj is None and schema is None:
             raise com.IbisError("One of the `schema` or `obj` parameter is required")
+        if schema is not None:
+            schema = ibis.schema(schema)
 
         if isinstance(obj, ir.Table) and schema is not None:
             if not schema.equals(obj.schema()):
diff --git a/ibis/backends/clickhouse/__init__.py b/ibis/backends/clickhouse/__init__.py
index 05016732727b6..f0ac9a78a0aae 100644
--- a/ibis/backends/clickhouse/__init__.py
+++ b/ibis/backends/clickhouse/__init__.py
@@ -666,13 +666,12 @@ def create_table(
 
         if obj is None and schema is None:
             raise com.IbisError("The `schema` or `obj` parameter is required")
+        if schema is not None:
+            schema = ibis.schema(schema)
 
         if obj is not None and not isinstance(obj, ir.Expr):
             obj = ibis.memtable(obj, schema=schema)
 
-        if schema is None:
-            schema = obj.schema()
-
         this = sge.Schema(
             this=sg.table(name, db=database),
             expressions=[
@@ -680,7 +679,7 @@ def create_table(
                     this=sg.to_identifier(name, quoted=self.compiler.quoted),
                     kind=self.compiler.type_mapper.from_ibis(typ),
                 )
-                for name, typ in schema.items()
+                for name, typ in (schema or obj.schema()).items()
             ],
         )
         properties = [
diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py
index f1e2f647b1569..13890d1b5d01c 100644
--- a/ibis/backends/duckdb/__init__.py
+++ b/ibis/backends/duckdb/__init__.py
@@ -41,6 +41,8 @@
     import torch
     from fsspec import AbstractFileSystem
 
+    from ibis.expr.schema import SchemaLike
+
 
 _UDF_INPUT_TYPE_MAPPING = {
     InputType.PYARROW: duckdb.functional.ARROW,
@@ -103,7 +105,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -147,6 +149,8 @@ def create_table(
 
         if obj is None and schema is None:
             raise ValueError("Either `obj` or `schema` must be specified")
+        if schema is not None:
+            schema = ibis.schema(schema)
 
         properties = []
 
diff --git a/ibis/backends/exasol/__init__.py b/ibis/backends/exasol/__init__.py
index 759dfc940d1d5..f72226627dccd 100644
--- a/ibis/backends/exasol/__init__.py
+++ b/ibis/backends/exasol/__init__.py
@@ -342,6 +342,8 @@ def create_table(
         """
         if obj is None and schema is None:
             raise ValueError("Either `obj` or `schema` must be specified")
+        if schema is not None:
+            schema = ibis.schema(schema)
 
         if temp:
             raise com.UnsupportedOperationError(
diff --git a/ibis/backends/mssql/__init__.py b/ibis/backends/mssql/__init__.py
index 508bbcf438308..cb03f61d9c3fc 100644
--- a/ibis/backends/mssql/__init__.py
+++ b/ibis/backends/mssql/__init__.py
@@ -605,6 +605,8 @@ def create_table(
         """
         if obj is None and schema is None:
             raise ValueError("Either `obj` or `schema` must be specified")
+        # if schema is not None:
+        #     schema = ibis.schema(schema)
 
         if temp and overwrite:
             raise ValueError(
diff --git a/ibis/backends/tests/test_client.py b/ibis/backends/tests/test_client.py
index 49848d8f43dc7..55dc7c732f53b 100644
--- a/ibis/backends/tests/test_client.py
+++ b/ibis/backends/tests/test_client.py
@@ -85,16 +85,11 @@ def _create_temp_table_with_schema(backend, con, temp_table_name, schema, data=N
     "sch",
     [
         None,
-        ibis.schema(
-            dict(
-                first_name="string",
-                last_name="string",
-                department_name="string",
-                salary="float64",
-            )
-        ),
+        dict(first_name="string", salary="float64"),
+        dict(first_name="string", salary="float64").items(),
+        ibis.schema(dict(first_name="string", salary="float64")),
     ],
-    ids=["no_schema", "schema"],
+    ids=["no_schema", "dict_schema", "tuples", "schema"],
 )
 @pytest.mark.notimpl(["druid"])
 @pytest.mark.notimpl(
@@ -102,14 +97,7 @@ def _create_temp_table_with_schema(backend, con, temp_table_name, schema, data=N
     reason="Flink backend supports creating only TEMPORARY VIEW for in-memory data.",
 )
 def test_create_table(backend, con, temp_table, func, sch):
-    df = pd.DataFrame(
-        {
-            "first_name": ["A", "B", "C"],
-            "last_name": ["D", "E", "F"],
-            "department_name": ["AA", "BB", "CC"],
-            "salary": [100.0, 200.0, 300.0],
-        }
-    )
+    df = pd.DataFrame({"first_name": ["A", "B", "C"], "salary": [100.0, 200.0, 300.0]})
 
     con.create_table(temp_table, func(df), schema=sch)
     result = (