Skip to content

Commit

Permalink
feat(trino): compile timestamp types with scale
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Jan 20, 2023
1 parent a38115a commit 67683d3
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 3 deletions.
4 changes: 3 additions & 1 deletion ibis/backends/base/sql/alchemy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -242,7 +242,9 @@ def _get_insert_method(self, expr):

def _columns_from_schema(self, name: str, schema: sch.Schema) -> list[sa.Column]:
return [
sa.Column(colname, to_sqla_type(dtype), nullable=dtype.nullable)
sa.Column(
colname, to_sqla_type(self.con.dialect, dtype), nullable=dtype.nullable
)
for colname, dtype in zip(schema.names, schema.types)
]

Expand Down
36 changes: 36 additions & 0 deletions ibis/backends/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import os
import platform
import re
import string

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -875,3 +876,38 @@ def test_get_backend(con, alltypes, monkeypatch):
assert ibis.get_backend() is con
expr = ibis.literal(1) + 2
assert ibis.get_backend(expr) is con


@pytest.mark.notyet(
[
"bigquery",
"dask",
"datafusion",
"duckdb",
"impala",
"mssql",
"mysql",
"pandas",
"polars",
"postgres",
"pyspark",
"sqlite",
],
reason="backend doesn't support timestamp with scale parameter",
)
@pytest.mark.notimpl(["clickhouse"], reason="create table isn't implemented")
@pytest.mark.notimpl(
["snowflake"], reason="scale not implemented in ibis's snowflake backend"
)
def test_create_table_timestamp(con):
schema = ibis.schema(
dict(zip(string.ascii_letters, map("timestamp({:d})".format, range(10))))
)
name = f"timestamp_scale_{guid()}"
con.create_table(name, schema=schema, force=True)
try:
rows = con.raw_sql(f"DESCRIBE {name}").fetchall()
result = ibis.schema((name, typ) for name, typ, *_ in rows)
assert result == schema
finally:
con.drop_table(name, force=True)
14 changes: 14 additions & 0 deletions ibis/backends/trino/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,20 @@ def _load_data(data_dir: Path, script_dir: Path, **_: Any) -> None:
"INSERT INTO map VALUES (MAP(ARRAY['d', 'e', 'f'], ARRAY[4, 5, 6]))"
)
)
c.execute(sa.text("DROP TABLE IF EXISTS ts"))
c.execute(
sa.text(
"CREATE TABLE ts (x TIMESTAMP(3), y TIMESTAMP(6), z TIMESTAMP(9))"
)
)
c.execute(
sa.text(
"INSERT INTO ts VALUES "
"(TIMESTAMP '2023-01-07 13:20:05.561', "
" TIMESTAMP '2023-01-07 13:20:05.561021', "
" TIMESTAMP '2023-01-07 13:20:05.561000231')"
)
)

@staticmethod
def connect(data_directory: Path):
Expand Down
11 changes: 9 additions & 2 deletions ibis/tests/benchmarks/test_benchmarks.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import numpy as np
import pandas as pd
import pytest
import sqlalchemy as sa
from packaging.version import parse as vparse

import ibis
Expand Down Expand Up @@ -188,7 +189,10 @@ def test_compile(benchmark, module, expr_fn, t, base, large_expr):
pytest.skip(str(e))
else:
expr = expr_fn(t, base, large_expr)
benchmark(mod.compile, expr)
try:
benchmark(mod.compile, expr)
except sa.exc.NoSuchModuleError as e:
pytest.skip(str(e))


@pytest.fixture(scope="module")
Expand Down Expand Up @@ -691,7 +695,10 @@ def test_compile_with_drops(
except (AttributeError, ImportError) as e:
pytest.skip(str(e))
else:
benchmark(mod.compile, expr)
try:
benchmark(mod.compile, expr)
except sa.exc.NoSuchModuleError as e:
pytest.skip(str(e))


def test_repr_join(benchmark, customers, orders, orders_items, products):
Expand Down

0 comments on commit 67683d3

Please sign in to comment.