feat(databricks): add support for json via variant

Author: cpcloud

ibis/backends/sql/compilers/databricks.py

@@ -5,6 +5,7 @@
 
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
+from ibis.backends.sql.compilers.base import NULL
 from ibis.backends.sql.compilers.pyspark import PySparkCompiler
 from ibis.backends.sql.dialects import Databricks
 
@@ -21,7 +22,14 @@ class DatabricksCompiler(PySparkCompiler):
         ops.BitXor: "bit_xor",
         ops.TypeOf: "typeof",
         ops.RandomUUID: "uuid",
+        ops.JSONGetItem: "json_extract",
     }
+    del (
+        SIMPLE_OPS[ops.UnwrapJSONString],
+        SIMPLE_OPS[ops.UnwrapJSONInt64],
+        SIMPLE_OPS[ops.UnwrapJSONFloat64],
+        SIMPLE_OPS[ops.UnwrapJSONBoolean],
+    )
 
     UNSUPPORTED_OPS = (
         ops.ElementWiseVectorizedUDF,
@@ -31,6 +39,42 @@ class DatabricksCompiler(PySparkCompiler):
         ops.TimestampBucket,
     )
 
+    def visit_ToJSONArray(self, op, *, arg):
+        return self.f.try_variant_get(arg, "$", "ARRAY<VARIANT>")
+
+    def visit_ToJSONMap(self, op, *, arg):
+        return self.f.try_variant_get(arg, "$", "MAP<STRING, VARIANT>")
+
+    def visit_UnwrapJSONString(self, op, *, arg):
+        return self.if_(
+            self.f.schema_of_variant(arg).eq(sge.convert("STRING")),
+            self.f.try_variant_get(arg, "$", "STRING"),
+            NULL,
+        )
+
+    def visit_UnwrapJSONInt64(self, op, *, arg):
+        return self.if_(
+            self.f.schema_of_variant(arg).eq(sge.convert("BIGINT")),
+            self.f.try_variant_get(arg, "$", "BIGINT"),
+            NULL,
+        )
+
+    def visit_UnwrapJSONFloat64(self, op, *, arg):
+        return self.if_(
+            self.f.schema_of_variant(arg).isin(
+                sge.convert("STRING"), sge.convert("BOOLEAN")
+            ),
+            NULL,
+            self.f.try_variant_get(arg, "$", "DOUBLE"),
+        )
+
+    def visit_UnwrapJSONBoolean(self, op, *, arg):
+        return self.if_(
+            self.f.schema_of_variant(arg).eq(sge.convert("BOOLEAN")),
+            self.f.try_variant_get(arg, "$", "BOOLEAN"),
+            NULL,
+        )
+
     def visit_NonNullLiteral(self, op, *, value, dtype):
         if dtype.is_binary():
             return self.f.unhex(value.hex())

ibis/backends/sql/datatypes.py

@@ -1325,6 +1325,16 @@ def _from_ibis_Map(cls, dtype: dt.Map) -> sge.DataType:
 class DatabricksType(SqlglotType):
     dialect = "databricks"
 
+    @classmethod
+    def _from_ibis_JSON(cls, dtype: dt.JSON) -> sge.DataType:
+        return sge.DataType(this=typecode.VARIANT)
+
+    @classmethod
+    def _from_sqlglot_VARIANT(cls, nullable: bool | None = None) -> sge.DataType:
+        return dt.JSON(nullable=nullable)
+
+    _from_sqlglot_JSON = _from_sqlglot_VARIANT
+
 
 class AthenaType(SqlglotType):
     dialect = "athena"

ibis/backends/tests/test_json.py

@@ -17,9 +17,7 @@
 pytestmark = [
     pytest.mark.never(["impala"], reason="doesn't support JSON and never will"),
     pytest.mark.notyet(["clickhouse"], reason="upstream is broken"),
-    pytest.mark.notimpl(
-        ["datafusion", "exasol", "mssql", "druid", "oracle", "databricks"]
-    ),
+    pytest.mark.notimpl(["datafusion", "exasol", "mssql", "druid", "oracle"]),
 ]