feat(api): support SchemaLike in Backend.create_table() (#9885)

Co-authored-by: Phillip Cloud <[email protected]>

ibis/backends/bigquery/__init__.py

@@ -923,7 +923,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -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()):

ibis/backends/clickhouse/__init__.py

@@ -611,7 +611,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -666,21 +666,20 @@ 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=[
                 sge.ColumnDef(
                     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 = [

ibis/backends/datafusion/__init__.py

@@ -596,7 +596,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: sch.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -625,6 +625,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 = []
 

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 = []
 

ibis/backends/exasol/__init__.py

@@ -314,7 +314,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: sch.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         overwrite: bool = False,
         temp: bool = False,
@@ -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(

ibis/backends/impala/__init__.py

@@ -463,7 +463,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema=None,
+        schema: sch.SchemaLike | None = None,
         database=None,
         temp: bool | None = None,
         overwrite: bool = False,
@@ -510,6 +510,8 @@ 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 temp is not None:
             raise NotImplementedError(
@@ -547,7 +549,7 @@ def create_table(
             self._safe_exec_sql(
                 CreateTableWithSchema(
                     name,
-                    schema if schema is not None else obj.schema(),
+                    schema or obj.schema(),
                     database=database or self.current_database,
                     format=format,
                     external=external,

ibis/backends/mssql/__init__.py

@@ -564,7 +564,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: sch.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -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(

ibis/backends/mysql/__init__.py

@@ -388,13 +388,15 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
     ) -> ir.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 = []
 

ibis/backends/oracle/__init__.py

@@ -374,7 +374,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -403,6 +403,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 = []
 

ibis/backends/pandas/__init__.py

@@ -210,6 +210,8 @@ 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:
             df = self._convert_object(obj)

ibis/backends/polars/__init__.py

@@ -362,7 +362,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool | None = None,
         overwrite: bool = False,

ibis/backends/postgres/__init__.py

@@ -626,7 +626,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -655,6 +655,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 = []
 

ibis/backends/pyspark/__init__.py

@@ -547,7 +547,7 @@ def create_table(
             ir.Table | pd.DataFrame | pa.Table | pl.DataFrame | pl.LazyFrame | None
         ) = None,
         *,
-        schema: sch.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool | None = None,
         overwrite: bool = False,
@@ -608,6 +608,7 @@ def create_table(
                 df = self._session.sql(query)
                 df.write.saveAsTable(name, format=format, mode=mode)
         elif schema is not None:
+            schema = ibis.schema(schema)
             schema = PySparkSchema.from_ibis(schema)
             with self._active_catalog_database(catalog, db):
                 self._session.catalog.createTable(name, schema=schema, format=format)

ibis/backends/risingwave/__init__.py

@@ -14,6 +14,7 @@
 import ibis.backends.sql.compilers as sc
 import ibis.common.exceptions as com
 import ibis.expr.operations as ops
+import ibis.expr.schema as sch
 import ibis.expr.types as ir
 from ibis import util
 from ibis.backends.postgres import Backend as PostgresBackend
@@ -130,7 +131,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -177,6 +178,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 connector_properties is not None and (
             encode_format is None or data_format is None

ibis/backends/snowflake/__init__.py

@@ -765,7 +765,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: sch.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -797,6 +797,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)
 
         quoted = self.compiler.quoted
 

ibis/backends/sqlite/__init__.py

@@ -436,7 +436,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: ibis.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,

ibis/backends/tests/test_client.py

@@ -85,31 +85,19 @@ 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(
     ["flink"],
     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 = (

ibis/backends/trino/__init__.py

@@ -395,7 +395,7 @@ def create_table(
         | pl.LazyFrame
         | None = None,
         *,
-        schema: sch.Schema | None = None,
+        schema: sch.SchemaLike | None = None,
         database: str | None = None,
         temp: bool = False,
         overwrite: bool = False,
@@ -435,6 +435,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 temp:
             raise NotImplementedError(