feat(alchemy): implement json getitem for sqlalchemy backends

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

@@ -256,7 +256,7 @@ def sa_inet(_, satype, nullable=True):
     return dt.INET(nullable=nullable)
 
 
-@dt.dtype.register(PGDialect, postgresql.JSON)
+@dt.dtype.register(Dialect, sa.types.JSON)
 def sa_json(_, satype, nullable=True):
     return dt.JSON(nullable=nullable)
 

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

@@ -179,6 +179,8 @@ def _cast(t, op):
         # in pandas. CAST(expr AS BYTEA) is correct and returns byte strings.
         return sa.cast(sa_arg, sa.LargeBinary())
 
+    if isinstance(typ, dt.JSON) and not t.native_json_type:
+        return sa_arg
     return sa.cast(sa_arg, sa_type)
 
 
@@ -646,6 +648,7 @@ def _count_star(t, op):
     ops.BitwiseLeftShift: _bitwise_op("<<"),
     ops.BitwiseRightShift: _bitwise_op(">>"),
     ops.BitwiseNot: _bitwise_not,
+    ops.JSONGetItem: fixed_arity(lambda x, y: x.op("->")(y), 2),
 }
 
 

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

@@ -47,6 +47,7 @@ class AlchemyExprTranslator(ExprTranslator):
     _has_reduction_filter_syntax = False
 
     integer_to_timestamp = sa.func.to_timestamp
+    native_json_type = True
 
     def name(self, translated, name, force=True):
         return translated.label(name)

ibis/backends/mysql/compiler.py

@@ -31,6 +31,7 @@ class MySQLExprTranslator(AlchemyExprTranslator):
     )
     _bool_aggs_need_cast_to_int32 = False
     integer_to_timestamp = sa.func.from_unixtime
+    native_json_type = False
 
 
 rewrites = MySQLExprTranslator.rewrites

ibis/backends/mysql/registry.py

@@ -169,6 +169,16 @@ def _find_in_set(t, op):
     )
 
 
+def _json_get_item(t, op):
+    arg = t.translate(op.arg)
+    index = t.translate(op.index)
+    if isinstance(op.index.output_dtype, dt.Integer):
+        path = "$[" + sa.cast(index, sa.TEXT) + "]"
+    else:
+        path = "$." + index
+    return sa.func.json_extract(arg, path)
+
+
 operation_registry.update(
     {
         ops.Literal: _literal,
@@ -211,5 +221,6 @@ def _find_in_set(t, op):
         ops.GroupConcat: _group_concat,
         ops.DayOfWeekIndex: _day_of_week_index,
         ops.DayOfWeekName: _day_of_week_name,
+        ops.JSONGetItem: _json_get_item,
     }
 )

ibis/backends/pandas/client.py

@@ -1,5 +1,6 @@
 """The pandas client implementation."""
 
+import json
 from collections.abc import Mapping, Sequence
 
 import numpy as np
@@ -302,6 +303,19 @@ def convert_array_to_series(in_dtype, out_dtype, column):
     return column.map(lambda x: x if x is None else list(x))
 
 
+@sch.convert.register(np.dtype, dt.JSON, pd.Series)
+def convert_json_to_series(in_, out, col: pd.Series):
+    def try_json(x):
+        if x is None:
+            return x
+        try:
+            return json.loads(x)
+        except (TypeError, json.JSONDecodeError):
+            return x
+
+    return pd.Series(list(map(try_json, col)), dtype="object")
+
+
 dt.DataType.to_pandas = ibis_dtype_to_pandas  # type: ignore
 sch.Schema.to_pandas = ibis_schema_to_pandas  # type: ignore