Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 3 commits into from
Aug 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion ibis/backends/bigquery/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -923,7 +923,7 @@
| pl.LazyFrame
| None = None,
*,
schema: ibis.Schema | None = None,
schema: sch.SchemaLike | None = None,
database: str | None = None,
temp: bool = False,
overwrite: bool = False,
Expand Down Expand Up @@ -972,6 +972,8 @@
"""
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
9 changes: 4 additions & 5 deletions ibis/backends/clickhouse/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down 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
4 changes: 3 additions & 1 deletion ibis/backends/datafusion/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 = []

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
4 changes: 3 additions & 1 deletion ibis/backends/exasol/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down 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
6 changes: 4 additions & 2 deletions ibis/backends/impala/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion ibis/backends/mssql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down 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
4 changes: 3 additions & 1 deletion ibis/backends/mysql/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []

Expand Down
4 changes: 3 additions & 1 deletion ibis/backends/oracle/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down 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: 1 addition & 1 deletion ibis/backends/polars/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion ibis/backends/postgres/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down 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
3 changes: 2 additions & 1 deletion ibis/backends/pyspark/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down 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
5 changes: 4 additions & 1 deletion ibis/backends/risingwave/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion ibis/backends/snowflake/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion ibis/backends/sqlite/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
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
4 changes: 3 additions & 1 deletion ibis/backends/trino/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down
Loading