refactor(datatype): introduce to_sqla_type dispatching on dialect

Author: cpcloud

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

@@ -402,14 +402,17 @@ def _handle_failed_column_type_inference(
         """Handle cases where SQLAlchemy cannot infer the column types of `table`."""
 
         self.inspector.reflect_table(table, table.columns)
-        quoted_name = self.con.dialect.identifier_preparer.quote(table.name)
+        dialect = self.con.dialect
+        quoted_name = dialect.identifier_preparer.quote(table.name)
 
         for colname, type in self._metadata(quoted_name):
             if colname in nulltype_cols:
                 # replace null types discovered by sqlalchemy with non null
                 # types
                 table.append_column(
-                    sa.Column(colname, to_sqla_type(type), nullable=type.nullable),
+                    sa.Column(
+                        colname, to_sqla_type(dialect, type), nullable=type.nullable
+                    ),
                     replace_existing=True,
                 )
         return table

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

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-import functools
 from typing import Iterable
 
 import sqlalchemy as sa
+from multipledispatch import Dispatcher
 from sqlalchemy.dialects import mssql, mysql, postgresql, sqlite
 from sqlalchemy.dialects.mssql.base import MSDialect
 from sqlalchemy.dialects.mysql.base import MySQLDialect
@@ -82,8 +82,9 @@ def table_from_schema(name, meta, schema, database: str | None = None):
     # Convert Ibis schema to SQLA table
     columns = []
 
+    dialect = getattr(meta.bind, "dialect", _DEFAULT_DIALECT)
     for colname, dtype in zip(schema.names, schema.types):
-        satype = to_sqla_type(dtype)
+        satype = to_sqla_type(dialect, dtype)
         column = sa.Column(colname, satype, nullable=dtype.nullable)
         columns.append(column)
 
@@ -115,55 +116,56 @@ def table_from_schema(name, meta, schema, database: str | None = None):
     dt.UInt32: UInt32,
     dt.UInt64: UInt64,
     dt.JSON: sa.JSON,
+    dt.Interval: sa.Interval,
 }
 
 
-@functools.singledispatch
-def to_sqla_type(itype, type_map=None):
-    if type_map is None:
-        type_map = ibis_type_to_sqla
-    return type_map[type(itype)]
+_DEFAULT_DIALECT = DefaultDialect()
 
+to_sqla_type = Dispatcher("to_sqla_type")
 
-@to_sqla_type.register(dt.Decimal)
-def _(itype, **kwargs):
-    return sa.types.NUMERIC(itype.precision, itype.scale)
+
+@to_sqla_type.register(Dialect, dt.DataType)
+def _default(_, itype):
+    return ibis_type_to_sqla[type(itype)]
 
 
-@to_sqla_type.register(dt.Interval)
-def _(itype, **kwargs):
-    return sa.types.Interval()
+@to_sqla_type.register(Dialect, dt.Decimal)
+def _decimal(_, itype):
+    return sa.types.NUMERIC(itype.precision, itype.scale)
 
 
-@to_sqla_type.register(dt.Date)
-def _(itype, **kwargs):
-    return sa.Date()
+@to_sqla_type.register(Dialect, dt.Timestamp)
+def _timestamp(_, itype):
+    return sa.TIMESTAMP(timezone=bool(itype.timezone))
 
 
-@to_sqla_type.register(dt.Timestamp)
-def _(itype, **kwargs):
-    return sa.TIMESTAMP(bool(itype.timezone))
+@to_sqla_type.register(Dialect, dt.Array)
+def _array(dialect, itype):
+    return ArrayType(to_sqla_type(dialect, itype.value_type))
 
 
-@to_sqla_type.register(dt.Array)
-def _(itype, **kwargs):
+@to_sqla_type.register(PGDialect, dt.Array)
+def _pg_array(dialect, itype):
     # Unwrap the array element type because sqlalchemy doesn't allow arrays of
     # arrays. This doesn't affect the underlying data.
     while itype.is_array():
         itype = itype.value_type
-    return sa.ARRAY(to_sqla_type(itype, **kwargs))
+    return sa.ARRAY(to_sqla_type(dialect, itype))
 
 
-@to_sqla_type.register(dt.Struct)
-def _(itype, **_):
+@to_sqla_type.register(Dialect, dt.Struct)
+def _struct(dialect, itype):
     return StructType(
-        [(name, to_sqla_type(type)) for name, type in itype.pairs.items()]
+        [(name, to_sqla_type(dialect, type)) for name, type in itype.pairs.items()]
     )
 
 
-@to_sqla_type.register(dt.Map)
-def _(itype, **_):
-    return MapType(to_sqla_type(itype.key_type), to_sqla_type(itype.value_type))
+@to_sqla_type.register(Dialect, dt.Map)
+def _map(dialect, itype):
+    return MapType(
+        to_sqla_type(dialect, itype.key_type), to_sqla_type(dialect, itype.value_type)
+    )
 
 
 @dt.dtype.register(Dialect, sa.types.NullType)
@@ -322,8 +324,8 @@ def ga_geometry(_, gatype, nullable=True):
         else:
             raise ValueError(f"Unrecognized geometry type: {t}")
 
-    @to_sqla_type.register(dt.GeoSpatial)
-    def _(itype, **kwargs):
+    @to_sqla_type.register(Dialect, dt.GeoSpatial)
+    def _(_, itype, **kwargs):
         if itype.geotype == 'geometry':
             return ga.Geometry
         elif itype.geotype == 'geography':
@@ -406,8 +408,8 @@ def sa_datetime(_, satype, nullable=True, default_timezone='UTC'):
     return dt.Timestamp(timezone=timezone, nullable=nullable)
 
 
-@dt.dtype.register(Dialect, sa.ARRAY)
-def sa_array(dialect, satype, nullable=True):
+@dt.dtype.register(PGDialect, sa.ARRAY)
+def sa_pg_array(dialect, satype, nullable=True):
     dimensions = satype.dimensions
     if dimensions is not None and dimensions != 1:
         raise NotImplementedError(

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

@@ -6,9 +6,7 @@
 from sqlalchemy import sql
 
 import ibis.expr.operations as ops
-import ibis.expr.schema as sch
 from ibis.backends.base.sql.alchemy.database import AlchemyTable
-from ibis.backends.base.sql.alchemy.datatypes import to_sqla_type
 from ibis.backends.base.sql.alchemy.translator import (
     AlchemyContext,
     AlchemyExprTranslator,
@@ -22,10 +20,6 @@
 from ibis.backends.base.sql.compiler.base import SetOp
 
 
-def _schema_to_sqlalchemy_columns(schema: sch.Schema) -> list[sa.Column]:
-    return [sa.column(n, to_sqla_type(t)) for n, t in schema.items()]
-
-
 class _AlchemyTableSetFormatter(TableSetFormatter):
     def get_result(self):
         # Got to unravel the join stack; the nesting order could be
@@ -94,32 +88,32 @@ def _format_table(self, op):
 
         alias = ctx.get_ref(op)
 
+        translator = ctx.compiler.translator_class(ref_op, ctx)
+
         if isinstance(ref_op, AlchemyTable):
             result = ref_op.sqla_table
         elif isinstance(ref_op, ops.UnboundTable):
             # use SQLAlchemy's TableClause for unbound tables
             result = sa.table(
-                ref_op.name,
-                *_schema_to_sqlalchemy_columns(ref_op.schema),
+                ref_op.name, *translator._schema_to_sqlalchemy_columns(ref_op.schema)
             )
         elif isinstance(ref_op, ops.SQLQueryResult):
-            columns = _schema_to_sqlalchemy_columns(ref_op.schema)
+            columns = translator._schema_to_sqlalchemy_columns(ref_op.schema)
             result = sa.text(ref_op.query).columns(*columns)
         elif isinstance(ref_op, ops.SQLStringView):
-            columns = _schema_to_sqlalchemy_columns(ref_op.schema)
+            columns = translator._schema_to_sqlalchemy_columns(ref_op.schema)
             result = sa.text(ref_op.query).columns(*columns).cte(ref_op.name)
         elif isinstance(ref_op, ops.View):
             # TODO(kszucs): avoid converting to expression
             child_expr = ref_op.child.to_expr()
             definition = child_expr.compile()
             result = sa.table(
-                ref_op.name,
-                *_schema_to_sqlalchemy_columns(ref_op.schema),
+                ref_op.name, *translator._schema_to_sqlalchemy_columns(ref_op.schema)
             )
             backend = child_expr._find_backend()
             backend._create_temp_view(view=result, definition=definition)
         elif isinstance(ref_op, ops.InMemoryTable):
-            columns = _schema_to_sqlalchemy_columns(ref_op.schema)
+            columns = translator._schema_to_sqlalchemy_columns(ref_op.schema)
 
             if self.context.compiler.cheap_in_memory_tables:
                 result = sa.table(ref_op.name, *columns)

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

@@ -5,7 +5,8 @@
 import ibis
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
-from ibis.backends.base.sql.alchemy.datatypes import ibis_type_to_sqla, to_sqla_type
+from ibis.backends.base.sql.alchemy import to_sqla_type
+from ibis.backends.base.sql.alchemy.datatypes import _DEFAULT_DIALECT
 from ibis.backends.base.sql.alchemy.registry import (
     fixed_arity,
     sqlalchemy_operation_registry,
@@ -35,7 +36,6 @@ def subcontext(self):
 class AlchemyExprTranslator(ExprTranslator):
     _registry = sqlalchemy_operation_registry
     _rewrites = ExprTranslator._rewrites.copy()
-    _type_map = ibis_type_to_sqla
 
     context_class = AlchemyContext
 
@@ -54,11 +54,26 @@ class AlchemyExprTranslator(ExprTranslator):
         ops.CumeDist,
     )
 
+    _dialect_name = "default"
+
+    @property
+    def dialect(self) -> sa.engine.interfaces.Dialect:
+        if (name := self._dialect_name) == "default":
+            return _DEFAULT_DIALECT
+        dialect_cls = sa.dialects.registry.load(name)
+        return dialect_cls()
+
+    def _schema_to_sqlalchemy_columns(self, schema):
+        return [
+            sa.column(name, to_sqla_type(self.dialect, dtype))
+            for name, dtype in schema.items()
+        ]
+
     def name(self, translated, name, force=True):
         return translated.label(name)
 
     def get_sqla_type(self, data_type):
-        return to_sqla_type(data_type, type_map=self._type_map)
+        return to_sqla_type(self.dialect, data_type)
 
     def _maybe_cast_bool(self, op, arg):
         if (

ibis/backends/clickhouse/datatypes.py

@@ -257,4 +257,9 @@ def _(ty: dt.Struct) -> str:
 
 @serialize_raw.register(dt.Timestamp)
 def _(ty: dt.Timestamp) -> str:
-    return "DateTime64(6)" if ty.timezone is None else f"DateTime64(6, {ty.timezone!r})"
+    if (scale := ty.scale) is None:
+        scale = 3
+
+    if (timezone := ty.timezone) is not None:
+        return f"DateTime64({scale:d}, {timezone})"
+    return f"DateTime64({scale:d})"

ibis/backends/clickhouse/tests/snapshots/test_functions/test_cast_string_col/timestamp/out.sql

@@ -1 +1 @@
-CAST(string_col AS Nullable(DateTime64(6)))
+CAST(string_col AS Nullable(DateTime64(3)))

ibis/backends/clickhouse/tests/snapshots/test_functions/test_timestamp_cast/out1.sql

@@ -1 +1 @@
-CAST(timestamp_col AS DateTime64(6))
+CAST(timestamp_col AS DateTime64(3))

ibis/backends/clickhouse/tests/snapshots/test_functions/test_timestamp_cast/out2.sql

@@ -1 +1 @@
-CAST(int_col AS DateTime64(6))
+CAST(int_col AS DateTime64(3))