fix(databricks): use AS JSON for programmatic output of schema information

Author: cpcloud

ibis/backends/databricks/__init__.py

@@ -5,6 +5,7 @@
 import contextlib
 import functools
 import getpass
+import json
 import os
 import sys
 import tempfile
@@ -20,23 +21,73 @@
 import ibis
 import ibis.backends.sql.compilers as sc
 import ibis.common.exceptions as exc
+import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 import ibis.expr.schema as sch
 import ibis.expr.types as ir
 from ibis import util
 from ibis.backends import CanCreateDatabase, PyArrowExampleLoader, UrlFromPath
 from ibis.backends.sql import SQLBackend
 from ibis.backends.sql.compilers.base import STAR, AlterTable, RenameTable
+from ibis.backends.sql.datatypes import DatabricksType
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
+    from collections.abc import Iterable, Mapping
 
     import pandas as pd
     import polars as pl
 
     from ibis.expr.schema import SchemaLike
 
 
+def _databricks_type_to_ibis(typ, nullable: bool = True) -> dt.DataType:
+    """Convert a Databricks type to an Ibis type."""
+    typname = typ["name"]
+    if typname == "array":
+        return dt.Array(
+            _databricks_type_to_ibis(
+                typ["element_type"], nullable=typ["element_nullable"]
+            ),
+            nullable=nullable,
+        )
+    elif typname == "map":
+        return dt.Map(
+            key_type=_databricks_type_to_ibis(typ["key_type"]),
+            value_type=_databricks_type_to_ibis(
+                typ["value_type"], nullable=typ["value_nullable"]
+            ),
+            nullable=nullable,
+        )
+    elif typname == "struct":
+        return dt.Struct(
+            {
+                field["name"]: _databricks_type_to_ibis(
+                    field["type"], nullable=field["nullable"]
+                )
+                for field in typ["fields"]
+            },
+            nullable=nullable,
+        )
+    elif typname == "decimal":
+        return dt.Decimal(
+            precision=typ["precision"], scale=typ["scale"], nullable=nullable
+        )
+    else:
+        return DatabricksType.from_string(typname, nullable=nullable)
+
+
+def _databricks_schema_to_ibis(schema: Iterable[Mapping[str, Any]]) -> sch.Schema:
+    """Convert a Databricks schema to an Ibis schema."""
+    return sch.Schema(
+        {
+            item["name"]: _databricks_type_to_ibis(
+                item["type"], nullable=item["nullable"]
+            )
+            for item in schema
+        }
+    )
+
+
 class Backend(SQLBackend, CanCreateDatabase, UrlFromPath, PyArrowExampleLoader):
     name = "databricks"
     compiler = sc.databricks.compiler
@@ -143,6 +194,9 @@ def create_table(
             else:
                 table = obj
 
+            if not schema:
+                schema = table.schema()
+
             self._run_pre_execute_hooks(table)
 
             query = self.compiler.to_sqlglot(table)
@@ -158,10 +212,7 @@ def create_table(
         dialect = self.dialect
 
         initial_table = sg.table(temp_name, catalog=catalog, db=database, quoted=quoted)
-        target = sge.Schema(
-            this=initial_table,
-            expressions=(schema or table.schema()).to_sqlglot(dialect),
-        )
+        target = sge.Schema(this=initial_table, expressions=schema.to_sqlglot(dialect))
 
         properties = sge.Properties(expressions=properties)
         create_stmt = sge.Create(kind="TABLE", this=target, properties=properties)
@@ -256,23 +307,18 @@ def get_schema(
         """
         table = sg.table(
             table_name, db=database, catalog=catalog, quoted=self.compiler.quoted
-        )
-        sql = sge.Describe(kind="TABLE", this=table).sql(self.dialect)
+        ).sql(self.dialect)
         try:
             with self.con.cursor() as cur:
-                out = cur.execute(sql).fetchall_arrow()
+                [(out,)] = cur.execute(f"DESCRIBE EXTENDED {table} AS JSON").fetchall()
         except databricks.sql.exc.ServerOperationError as e:
             raise exc.TableNotFound(
                 f"Table {table_name!r} not found in "
                 f"{catalog or self.current_catalog}.{database or self.current_database}"
             ) from e
 
-        names = out["col_name"].to_pylist()
-        types = out["data_type"].to_pylist()
-
-        return sch.Schema(
-            dict(zip(names, map(self.compiler.type_mapper.from_string, types)))
-        )
+        js = json.loads(out)
+        return _databricks_schema_to_ibis(js["columns"])
 
     @contextlib.contextmanager
     def _safe_raw_sql(self, query, *args, **kwargs):

ibis/backends/databricks/tests/conftest.py

@@ -6,6 +6,8 @@
 from os import environ as env
 from typing import TYPE_CHECKING, Any
 
+import pytest
+
 import ibis
 from ibis.backends.tests.base import BackendTest
 
@@ -73,3 +75,9 @@ def connect(*, tmpdir, worker_id, **kw) -> BaseBackend:
             schema="default",
             **kw,
         )
+
+
+@pytest.fixture(scope="session")
+def con(tmp_path_factory, data_dir, worker_id):
+    with TestConf.load_data(data_dir, tmp_path_factory, worker_id) as be:
+        yield be.connection

ibis/backends/databricks/tests/test_datatypes.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import pytest
+
+import ibis.expr.datatypes as dt
+import ibis.expr.schema as sch
+from ibis.util import gen_name
+
+
+@pytest.fixture
+def tmp_table(con):
+    name = gen_name("databricks_tmp_table")
+    yield name
+    con.drop_table(name, force=True)
+
+
+def test_nested(con, tmp_table):
+    schema = sch.Schema(
+        {
+            "nums": "decimal(33, 2)",
+            "time": "timestamp('UTC')",
+            "operationName": "string",
+            "category": "string",
+            "tenantId": "string",
+            "properties": dt.Struct(
+                {
+                    "Timestamp": "timestamp('UTC')",
+                    "ActionType": "string",
+                    "Application": "string",
+                    "ApplicationId": "int32",
+                    "AppInstanceId": "int32",
+                    "AccountObjectId": "string",
+                    "AccountId": "string",
+                    "AccountDisplayName": "string",
+                    "IsAdminOperation": "boolean",
+                    "DeviceType": "string",
+                    "OSPlatform": "string",
+                    "IPAddress": "string",
+                    "IsAnonymousProxy": "boolean",
+                    "CountryCode": "string",
+                    "City": "string",
+                    "ISP": "string",
+                    "UserAgent": "string",
+                    "ActivityType": "string",
+                    "ActivityObjects": "string",
+                    "ObjectName": "string",
+                    "ObjectType": "string",
+                    "ObjectId": "string",
+                    "ReportId": "string",
+                    "AccountType": "string",
+                    "IsExternalUser": "boolean",
+                    "IsImpersonated": "boolean",
+                    "IPTags": "string",
+                    "IPCategory": "string",
+                    "UserAgentTags": "string",
+                    "RawEventData": "string",
+                    "AdditionalFields": "string",
+                }
+            ),
+            "Tenant": "string",
+            "_rescued_data": "string",
+            "timestamp": "timestamp('UTC')",
+            "parse_details": dt.Struct(
+                {
+                    "status": "string",
+                    "at": "timestamp('UTC')",
+                    "info": dt.Struct({"input-file-name": "string"}),
+                }
+            ),
+            "p_date": "string",
+            "foo": "array<struct<a: map<string, array<struct<b: string>>>>>",
+            "a": "array<int>",
+            "b": "map<string, int>",
+        }
+    )
+
+    t = con.create_table(tmp_table, schema=schema)
+    assert t.schema() == schema

ibis/backends/tests/test_client.py

@@ -407,7 +407,7 @@ def test_rename_table(con, temp_table, temp_table_orig):
 
 
 @mark.notimpl(["polars", "druid", "athena"])
-@mark.never(["impala", "pyspark", "databricks"], reason="No non-nullable datatypes")
+@mark.never(["impala", "pyspark"], reason="No non-nullable datatypes")
 @pytest.mark.notimpl(
     ["flink"],
     raises=com.IbisError,