Skip to content

Commit

Permalink
fix(trino): ensure that nested array types are inferred correctly
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed May 12, 2023
1 parent 012557a commit 030f76d
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 6 deletions.
7 changes: 1 addition & 6 deletions ibis/backends/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ def test_array_index(con, idx):

@builtin_array
@pytest.mark.never(
["clickhouse", "duckdb", "pandas", "pyspark", "snowflake", "polars"],
["clickhouse", "duckdb", "pandas", "pyspark", "snowflake", "polars", "trino"],
reason="backend does not flatten array types",
raises=AssertionError,
)
Expand Down Expand Up @@ -255,11 +255,6 @@ def test_array_discovery_clickhouse(backend):
reason="backend does not support nullable nested types",
raises=AssertionError,
)
@pytest.mark.notimpl(
["trino"],
reason="trino supports nested arrays, but not with the postgres connector",
raises=AssertionError,
)
@pytest.mark.never(
["bigquery"],
reason="doesn't support arrays of arrays",
Expand Down
5 changes: 5 additions & 0 deletions ibis/backends/trino/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import ibis.expr.datatypes as dt
from ibis import util
from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
from ibis.backends.base.sql.alchemy.datatypes import ArrayType
from ibis.backends.trino.compiler import TrinoSQLCompiler
from ibis.backends.trino.datatypes import ROW, parse

Expand Down Expand Up @@ -85,6 +86,10 @@ def _new_sa_metadata():
def column_reflect(inspector, table, column_info):
if isinstance(typ := column_info["type"], _ROW):
column_info["type"] = ROW(typ.attr_types)
elif isinstance(typ, sa.ARRAY):
column_info["type"] = toolz.nth(
typ.dimensions or 1, toolz.iterate(ArrayType, typ.item_type)
)

return meta

Expand Down

0 comments on commit 030f76d

Please sign in to comment.