Skip to content

Commit

Permalink
feat: support SchemaLike in Backend.create_table()
Browse files Browse the repository at this point in the history
  • Loading branch information
NickCrews committed Aug 24, 2024
1 parent 912a8d1 commit ccb001a
Show file tree
Hide file tree
Showing 12 changed files with 30 additions and 22 deletions.
2 changes: 2 additions & 0 deletions ibis/backends/bigquery/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Check warning on line 976 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L976

Added line #L976 was not covered by tests

if isinstance(obj, ir.Table) and schema is not None:
if not schema.equals(obj.schema()):
Expand Down
7 changes: 3 additions & 4 deletions ibis/backends/clickhouse/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
6 changes: 5 additions & 1 deletion ibis/backends/duckdb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 = []

Expand Down
2 changes: 2 additions & 0 deletions ibis/backends/exasol/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
2 changes: 2 additions & 0 deletions ibis/backends/mssql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
2 changes: 2 additions & 0 deletions ibis/backends/mysql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,8 @@ def create_table(
) -> 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 = []

Expand Down
2 changes: 2 additions & 0 deletions ibis/backends/oracle/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []

Expand Down
2 changes: 2 additions & 0 deletions ibis/backends/pandas/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions ibis/backends/postgres/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []

Expand Down
1 change: 1 addition & 0 deletions ibis/backends/pyspark/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 2 additions & 0 deletions ibis/backends/risingwave/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,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
Expand Down
22 changes: 5 additions & 17 deletions ibis/backends/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand Down

0 comments on commit ccb001a

Please sign in to comment.