refactor(datatype): add custom sqlalchemy nested types for backend differentiation

Author: cpcloud

ibis/backends/base/sql/alchemy/datatypes.py

@@ -22,25 +22,43 @@
     import geoalchemy2 as ga
 
 
+class ArrayType(UserDefinedType):
+    def __init__(self, value_type: sa.types.TypeEngine):
+        self.value_type = sa.types.to_instance(value_type)
+
+
+@compiles(ArrayType, "default")
+def compiles_array(element, compiler, **kw):
+    return f"ARRAY({compiler.process(element.value_type, **kw)})"
+
+
 class StructType(UserDefinedType):
     def __init__(
         self,
         pairs: Iterable[tuple[str, sa.types.TypeEngine]],
     ):
         self.pairs = [(name, sa.types.to_instance(type)) for name, type in pairs]
 
-    def get_col_spec(self, **_):
-        pairs = ", ".join(f"{k} {v}" for k, v in self.pairs)
-        return f"STRUCT({pairs})"
+
+@compiles(StructType, "default")
+def compiles_struct(element, compiler, **kw):
+    content = ", ".join(
+        f"{field} {compiler.process(typ, **kw)}" for field, typ in element.pairs
+    )
+    return f"STRUCT({content})"
 
 
 class MapType(UserDefinedType):
     def __init__(self, key_type: sa.types.TypeEngine, value_type: sa.types.TypeEngine):
         self.key_type = sa.types.to_instance(key_type)
         self.value_type = sa.types.to_instance(value_type)
 
-    def get_col_spec(self, **_):
-        return f"MAP({self.key_type}, {self.value_type})"
+
+@compiles(MapType, "default")
+def compiles_map(element, compiler, **kw):
+    key_type = compiler.process(element.key_type, **kw)
+    value_type = compiler.process(element.value_type, **kw)
+    return f"MAP({key_type}, {value_type})"
 
 
 class UInt64(sa.types.Integer):
@@ -426,6 +444,11 @@ def sa_struct(dialect, satype, nullable=True):
     return dt.Struct.from_tuples(pairs, nullable=nullable)
 
 
+@dt.dtype.register(Dialect, ArrayType)
+def sa_array(dialect, satype, nullable=True):
+    return dt.Array(dt.dtype(dialect, satype.value_type), nullable=nullable)
+
+
 @sch.infer.register((sa.Table, sa.sql.TableClause))
 def schema_from_table(
     table: sa.Table,

ibis/backends/base/sql/alchemy/registry.py

@@ -185,7 +185,7 @@ def _cast(t, op):
     if typ.is_json() and not t.native_json_type:
         return sa_arg
 
-    return t.cast(sa_arg, typ)
+    return sa.cast(sa_arg, t.get_sqla_type(typ))
 
 
 def _contains(func):

ibis/backends/base/sql/alchemy/translator.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import functools
+
 import sqlalchemy as sa
 
 import ibis
@@ -56,7 +58,7 @@ class AlchemyExprTranslator(ExprTranslator):
 
     _dialect_name = "default"
 
-    @property
+    @functools.cached_property
     def dialect(self) -> sa.engine.interfaces.Dialect:
         if (name := self._dialect_name) == "default":
             return _DEFAULT_DIALECT
@@ -65,8 +67,7 @@ def dialect(self) -> sa.engine.interfaces.Dialect:
 
     def _schema_to_sqlalchemy_columns(self, schema):
         return [
-            sa.column(name, to_sqla_type(self.dialect, dtype))
-            for name, dtype in schema.items()
+            sa.column(name, self.get_sqla_type(dtype)) for name, dtype in schema.items()
         ]
 
     def name(self, translated, name, force=True):
@@ -103,9 +104,6 @@ def _reduction(self, sa_func, op):
 
         return sa_func(*sa_args)
 
-    def cast(self, sa_expr, ibis_type: dt.DataType):
-        return sa.cast(sa_expr, self.get_sqla_type(ibis_type))
-
 
 rewrites = AlchemyExprTranslator.rewrites
 

ibis/backends/duckdb/compiler.py

@@ -5,7 +5,11 @@
 import ibis.backends.base.sql.alchemy.datatypes as sat
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
-from ibis.backends.base.sql.alchemy import AlchemyCompiler, AlchemyExprTranslator
+from ibis.backends.base.sql.alchemy import (
+    AlchemyCompiler,
+    AlchemyExprTranslator,
+    to_sqla_type,
+)
 from ibis.backends.duckdb.registry import operation_registry
 
 
@@ -24,6 +28,11 @@ def compile_uint(element, compiler, **kw):
     return element.__class__.__name__.upper()
 
 
+@compiles(sat.ArrayType, "duckdb")
+def compile_array(element, compiler, **kw):
+    return f"{compiler.process(element.value_type, **kw)}[]"
+
+
 try:
     import duckdb_engine
 except ImportError:
@@ -37,6 +46,14 @@ def compile_uint(element, compiler, **kw):
     def dtype_uint(_, satype, nullable=True):
         return getattr(dt, satype.__class__.__name__)(nullable=nullable)
 
+    @dt.dtype.register(duckdb_engine.Dialect, sat.ArrayType)
+    def _(dialect, satype, nullable=True):
+        return dt.Array(dt.dtype(dialect, satype.value_type), nullable=nullable)
+
+    @to_sqla_type.register(duckdb_engine.Dialect, dt.Array)
+    def _(dialect, itype):
+        return sat.ArrayType(to_sqla_type(dialect, itype.value_type))
+
 
 rewrites = DuckDBSQLExprTranslator.rewrites
 

ibis/backends/snowflake/__init__.py

@@ -41,11 +41,6 @@ class SnowflakeExprTranslator(AlchemyExprTranslator):
     _require_order_by = (*AlchemyExprTranslator._require_order_by, ops.Reduction)
     _dialect_name = "snowflake"
 
-    def cast(self, sa_expr, ibis_type: dt.DataType):
-        if ibis_type.is_array() or ibis_type.is_map() or ibis_type.is_struct():
-            return sa.type_coerce(sa_expr, self.get_sqla_type(ibis_type))
-        return super().cast(sa_expr, ibis_type)
-
 
 class SnowflakeCompiler(AlchemyCompiler):
     translator_class = SnowflakeExprTranslator

ibis/backends/snowflake/datatypes.py

@@ -15,6 +15,7 @@
 from snowflake.sqlalchemy.snowdialect import SnowflakeDialect
 
 import ibis.expr.datatypes as dt
+from ibis.backends.base.sql.alchemy import to_sqla_type
 
 if TYPE_CHECKING:
     from ibis.expr.datatypes import DataType
@@ -177,3 +178,18 @@ def sa_sf_numeric(_, satype, nullable=True):
 @dt.dtype.register(SnowflakeDialect, (sa.REAL, sa.FLOAT, sa.Float))
 def sa_sf_real_float(_, satype, nullable=True):
     return dt.Float64(nullable=nullable)
+
+
+@to_sqla_type.register(SnowflakeDialect, dt.Array)
+def _sf_array(_, itype):
+    return ARRAY
+
+
+@to_sqla_type.register(SnowflakeDialect, (dt.Map, dt.Struct))
+def _sf_map_struct(_, itype):
+    return OBJECT
+
+
+@to_sqla_type.register(SnowflakeDialect, dt.JSON)
+def _sf_json(_, itype):
+    return VARIANT

ibis/backends/tests/test_array.py

@@ -150,11 +150,11 @@ def test_array_index(con, idx):
 
 @builtin_array
 @pytest.mark.never(
-    ["clickhouse", "pandas", "pyspark", "snowflake", "polars"],
+    ["clickhouse", "duckdb", "pandas", "pyspark", "snowflake", "polars"],
     reason="backend does not flatten array types",
 )
 @pytest.mark.never(["bigquery"], reason="doesn't support arrays of arrays")
-def test_array_discovery_postgres_duckdb(con):
+def test_array_discovery_postgres(con):
     t = con.table("array_types")
     expected = ibis.schema(
         dict(
@@ -195,8 +195,7 @@ def test_array_discovery_clickhouse(con):
 
 @builtin_array
 @pytest.mark.notyet(
-    ["clickhouse", "duckdb", "postgres"],
-    reason="backend does not support nullable nested types",
+    ["clickhouse", "postgres"], reason="backend does not support nullable nested types"
 )
 @pytest.mark.notimpl(
     ["trino"],

ibis/backends/trino/datatypes.py

@@ -6,8 +6,8 @@
 from trino.sqlalchemy.datatype import DOUBLE, JSON, MAP, ROW
 from trino.sqlalchemy.dialect import TrinoDialect
 
-import ibis.backends.base.sql.alchemy.datatypes as sat
 import ibis.expr.datatypes as dt
+from ibis.backends.base.sql.alchemy import to_sqla_type
 from ibis.common.parsing import (
     COMMA,
     FIELD,
@@ -192,7 +192,8 @@ def _compiles_row(element, compiler, **kw):
     # TODO: @compiles should live in the dialect
     quote = compiler.dialect.identifier_preparer.quote
     content = ", ".join(
-        f"{field} {compiler.process(typ, **kw)}" for field, typ in element.pairs
+        f"{quote(field)} {compiler.process(typ, **kw)}"
+        for field, typ in element.attr_types
     )
     return f"ROW({content})"
 
@@ -206,6 +207,7 @@ def _map(dialect, itype):
 
 @compiles(MAP, "trino")
 def compiles_map(typ, compiler, **kw):
+    # TODO: @compiles should live in the dialect
     key_type = compiler.process(typ.key_type, **kw)
     value_type = compiler.process(typ.value_type, **kw)
     return f"MAP({key_type}, {value_type})"

ibis/tests/sql/test_sqlalchemy.py

@@ -29,6 +29,7 @@
     schema_from_table,
     to_sqla_type,
 )
+from ibis.backends.base.sql.alchemy.datatypes import ArrayType
 from ibis.tests.expr.mocks import MockAlchemyBackend
 from ibis.tests.util import assert_decompile_roundtrip, assert_equal
 
@@ -1107,7 +1108,7 @@ def test_to_sqla_type_array_of_non_primitive():
     expected_type = sa.BigInteger()
     assert result_name == expected_name
     assert type(result_type) == type(expected_type)
-    assert isinstance(result, sa.ARRAY)
+    assert isinstance(result, ArrayType)
 
 
 def test_no_cart_join(con, snapshot):