feat(dtype): support compiling dtypes to sql

Author: NickCrews

ibis/backends/sql/compilers/base.py

@@ -7,7 +7,7 @@
 import operator
 import string
 from functools import partial, reduce
-from typing import TYPE_CHECKING, Any, Callable, ClassVar
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, overload
 
 import sqlglot as sg
 import sqlglot.expressions as sge
@@ -17,6 +17,7 @@
 import ibis.common.patterns as pats
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
+import ibis.expr.schema as sch
 from ibis.backends.sql.rewrites import (
     FirstValue,
     LastValue,
@@ -53,7 +54,6 @@ def AlterTable(*args, kind="TABLE", **kwargs):
 if TYPE_CHECKING:
     from collections.abc import Iterable, Mapping
 
-    import ibis.expr.schema as sch
     import ibis.expr.types as ir
     from ibis.backends.sql.datatypes import SqlglotType
 
@@ -575,16 +575,52 @@ def _prepare_params(self, params):
             result[node] = value
         return result
 
+    @overload
     def to_sqlglot(
         self,
-        expr: ir.Expr,
+        x: dt.DataType,
+        *,
+        params: Mapping[ir.Expr, Any] | None = None,
+    ) -> sge.DataType: ...
+
+    @overload
+    def to_sqlglot(
+        self,
+        x: sch.Schema,
+        *,
+        params: Mapping[ir.Expr, Any] | None = None,
+    ) -> list[sge.ColumnDef]: ...
+
+    @overload
+    def to_sqlglot(
+        self,
+        x: ir.Expr,
         *,
         limit: str | None = None,
         params: Mapping[ir.Expr, Any] | None = None,
-    ):
+    ) -> sge.Expression: ...
+
+    def to_sqlglot(
+        self,
+        x: ir.Expr | dt.DataType | sch.Schema,
+        *,
+        limit: str | None = None,
+        params: Mapping[ir.Expr, Any] | None = None,
+    ) -> sge.Expression:
+        if isinstance(x, dt.DataType | sch.Schema):
+            return x.to_sqlglot(self.dialect)
+        return self._to_sqlglot_expr(x, limit=limit, params=params)
+
+    def _to_sqlglot_expr(
+        self,
+        x: ir.Expr,
+        *,
+        limit: str | None = None,
+        params: Mapping[ir.Expr, Any] | None = None,
+    ) -> sge.Expression:
         import ibis
 
-        table_expr = expr.as_table()
+        table_expr = x.as_table()
 
         if limit == "default":
             limit = ibis.options.sql.default_limit

ibis/backends/sql/compilers/bigquery/__init__.py

@@ -14,6 +14,7 @@
 import ibis.common.exceptions as com
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
+import ibis.expr.schema as sch
 from ibis import util
 from ibis.backends.sql.compilers.base import NULL, STAR, AggGen, SQLGlotCompiler
 from ibis.backends.sql.compilers.bigquery.udf.core import PythonToJavaScriptTranslator
@@ -206,7 +207,7 @@ class BigQueryCompiler(SQLGlotCompiler):
 
     def to_sqlglot(
         self,
-        expr: ir.Expr,
+        expr: ir.Expr | dt.DataType | sch.Schema,
         *,
         limit: str | None = None,
         params: Mapping[ir.Expr, Any] | None = None,
@@ -234,17 +235,18 @@ def to_sqlglot(
         Any
             The output of compilation. The type of this value depends on the
             backend.
-
         """
-        sql = super().to_sqlglot(expr, limit=limit, params=params)
+        if isinstance(expr, dt.DataType | sch.Schema):
+            return expr.to_sqlglot(self.dialect)
+        sgexpr = super().to_sqlglot(expr, limit=limit, params=params)
 
         table_expr = expr.as_table()
 
         memtable_names = frozenset(
             op.name for op in table_expr.op().find(ops.InMemoryTable)
         )
 
-        result = sql.transform(
+        result = sgexpr.transform(
             _qualify_memtable,
             dataset=session_dataset_id,
             project=session_project,
@@ -257,8 +259,8 @@ def to_sqlglot(
             compile_func = getattr(
                 self, f"_compile_{udf_node.__input_type__.name.lower()}_udf"
             )
-            if sql := compile_func(udf_node):
-                sources.append(sql)
+            if sgexpr := compile_func(udf_node):
+                sources.append(sgexpr)
 
         if not sources:
             return result

ibis/backends/sql/compilers/duckdb.py

@@ -107,20 +107,20 @@ class DuckDBCompiler(SQLGlotCompiler):
         ops.RandomScalar: "random",
     }
 
-    def to_sqlglot(
+    def _to_sqlglot_expr(
         self,
         expr: ir.Expr,
         *,
         limit: str | None = None,
         params: Mapping[ir.Expr, Any] | None = None,
     ):
-        sql = super().to_sqlglot(expr, limit=limit, params=params)
+        sgexpr = super()._to_sqlglot_expr(expr, limit=limit, params=params)
 
         table_expr = expr.as_table()
         geocols = table_expr.schema().geospatial
 
         if not geocols:
-            return sql
+            return sgexpr
 
         quoted = self.quoted
         return sg.select(
@@ -132,7 +132,7 @@ def to_sqlglot(
                     for col in geocols
                 ]
             )
-        ).from_(sql.subquery())
+        ).from_(sgexpr.subquery())
 
     def visit_StructColumn(self, op, *, names, values):
         return sge.Struct.from_arg_list(

ibis/backends/sql/compilers/mssql.py

@@ -35,7 +35,7 @@
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
-    import ibis.expr.operations as ir
+    from ibis.expr import types as ir
 
 y = var("y")
 start = var("start")
@@ -156,14 +156,13 @@ def _minimize_spec(op, spec):
             return None
         return spec
 
-    def to_sqlglot(
+    def _to_sqlglot_expr(
         self,
         expr: ir.Expr,
         *,
         limit: str | None = None,
         params: Mapping[ir.Expr, Any] | None = None,
-    ):
-        """Compile an Ibis expression to a sqlglot object."""
+    ) -> sge.Expression:
         import ibis
 
         table_expr = expr.as_table()
@@ -175,7 +174,7 @@ def to_sqlglot(
 
         if conversions:
             table_expr = table_expr.mutate(**conversions)
-        return super().to_sqlglot(table_expr, limit=limit, params=params)
+        return super()._to_sqlglot_expr(table_expr, limit=limit, params=params)
 
     def visit_RandomScalar(self, op):
         # By default RAND() will generate the same value for all calls within a

ibis/backends/sql/compilers/postgres.py

@@ -122,20 +122,20 @@ class PostgresCompiler(SQLGlotCompiler):
         ops.RandomUUID: "gen_random_uuid",
     }
 
-    def to_sqlglot(
+    def _to_sqlglot_expr(
         self,
         expr: ir.Expr,
         *,
         limit: str | None = None,
         params: Mapping[ir.Expr, Any] | None = None,
-    ):
+    ) -> sg.Expression:
         table_expr = expr.as_table()
         geocols = table_expr.schema().geospatial
         conversions = {name: table_expr[name].as_ewkb() for name in geocols}
 
         if conversions:
             table_expr = table_expr.mutate(**conversions)
-        return super().to_sqlglot(table_expr, limit=limit, params=params)
+        return super()._to_sqlglot_expr(table_expr, limit=limit, params=params)
 
     def _compile_python_udf(self, udf_node: ops.ScalarUDF):
         config = udf_node.__config__

ibis/backends/sql/datatypes.py

@@ -172,11 +172,12 @@ def to_ibis(cls, typ: sge.DataType, nullable: bool | None = None) -> dt.DataType
     @classmethod
     def from_ibis(cls, dtype: dt.DataType) -> sge.DataType:
         """Convert an Ibis dtype to an sqlglot dtype."""
-
         if method := getattr(cls, f"_from_ibis_{dtype.name}", None):
             return method(dtype)
         else:
-            return sge.DataType(this=_to_sqlglot_types[type(dtype)])
+            return sge.DataType(
+                this=_to_sqlglot_types[type(dtype)], nullable=dtype.nullable
+            )
 
     @classmethod
     def from_string(cls, text: str, nullable: bool | None = None) -> dt.DataType:
@@ -1346,7 +1347,15 @@ class AthenaType(SqlglotType):
     dialect = "athena"
 
 
-TYPE_MAPPERS = {
+_TYPE_MAPPERS = {
     mapper.dialect: mapper
     for mapper in set(get_subclasses(SqlglotType)) - {SqlglotType, BigQueryUDFType}
 }
+_TYPE_MAPPERS["pyspark"] = PySparkType
+_TYPE_MAPPERS["druid"] = DruidType
+
+
+def get_mapper(dialect: str | type[sg.Dialect]) -> type[SqlglotType]:
+    if not isinstance(dialect, str):
+        dialect = dialect.__name__.lower()
+    return _TYPE_MAPPERS[dialect]

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/athena/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/bigquery/to_sql_dtype.sql

@@ -0,0 +1 @@
+STRING

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/clickhouse/to_sql_dtype.sql

@@ -0,0 +1 @@
+Nullable(String)

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/databricks/to_sql_dtype.sql

@@ -0,0 +1 @@
+STRING

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/datafusion/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/druid/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/duckdb/to_sql_dtype.sql

@@ -0,0 +1 @@
+TEXT

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/exasol/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR(2000000)

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/flink/to_sql_dtype.sql

@@ -0,0 +1 @@
+STRING

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/impala/to_sql_dtype.sql

@@ -0,0 +1 @@
+STRING

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/mssql/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR(max)

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/mysql/to_sql_dtype.sql

@@ -0,0 +1 @@
+TEXT

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/oracle/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR2(4000)

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/postgres/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/pyspark/to_sql_dtype.sql

@@ -0,0 +1 @@
+STRING

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/risingwave/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/snowflake/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/sqlite/to_sql_dtype.sql

@@ -0,0 +1 @@
+TEXT

ibis/backends/tests/snapshots/test_sql/test_to_sql_dtype/trino/to_sql_dtype.sql

@@ -0,0 +1 @@
+VARCHAR

ibis/backends/tests/test_sql.py

@@ -197,6 +197,33 @@ def test_to_sql_default_backend(con, snapshot, monkeypatch):
     snapshot.assert_match(ibis.to_sql(expr), "to_sql.sql")
 
 
+@pytest.mark.parametrize(
+    "dialect",
+    [
+        # Just check these two to make sure that everything is plumbed through
+        pytest.param("sqlite", marks=pytest.mark.xfail(reason="arrays not supported")),
+        "duckdb",
+    ],
+)
+def test_to_sql_dtype_default_backend(dialect):
+    dt = ibis.dtype("array<int64>")
+    original_backend = ibis.get_backend()
+    try:
+        ibis.set_backend(dialect)
+        sql = ibis.to_sql(dt)
+    except:
+        ibis.set_backend(original_backend)
+        raise
+    assert "BIGINT[]" == sql
+
+
+@pytest.mark.parametrize("backend_name", _get_backends_to_test(discard=("polars",)))
+def test_to_sql_dtype(backend_name, snapshot):
+    dt = ibis.dtype("string")
+    sql = ibis.to_sql(dt, dialect=backend_name)
+    snapshot.assert_match(sql, "to_sql_dtype.sql")
+
+
 @pytest.mark.notimpl(["polars"], raises=ValueError, reason="not a SQL backend")
 def test_many_subqueries(backend_name, snapshot):
     def query(t, group_cols):

ibis/expr/datatypes/core.py

@@ -8,6 +8,7 @@
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from numbers import Integral, Real
 from typing import (
+    TYPE_CHECKING,
     Annotated,
     Any,
     Generic,
@@ -29,6 +30,10 @@
 from ibis.common.patterns import Between, Coercible, CoercionError
 from ibis.common.temporal import IntervalUnit, TimestampUnit
 
+if TYPE_CHECKING:
+    import sqlglot as sg
+    import sqlglot.expressions as sge
+
 
 @lazy_singledispatch
 def dtype(value: Any, nullable: bool = True) -> DataType:
@@ -263,6 +268,33 @@ def from_polars(cls, polars_type, nullable=True) -> Self:
 
         return PolarsType.to_ibis(polars_type, nullable=nullable)
 
+    def to_sqlglot(self, dialect: str | sg.Dialect) -> sge.DataType:
+        """Convert to the equivalent sqlglot.DataType.
+
+        Parameters
+        ----------
+        dialect
+            The SQL dialect to use.
+            For example, some dialects convert an ibis string to TEXT,
+            while others use VARCHAR.
+
+        Returns
+        -------
+        DataType
+            The equivalent sqlglot.DataType.
+
+        Examples
+        --------
+        >>> import ibis
+        >>> dt = ibis.dtype("!string")
+        >>> dt.to_sqlglot(dialect="duckdb")
+        DataType(this=Type.VARCHAR)
+        """
+        from ibis.backends.sql.datatypes import get_mapper
+
+        type_mapper = get_mapper(dialect)
+        return type_mapper.from_ibis(self)
+
     def to_numpy(self):
         """Return the equivalent numpy datatype."""
         from ibis.formats.numpy import NumpyType