feat(trino): implement json getitem support

ibis-project · Dec 21, 2022 · 7c41566 · 7c41566
1 parent 0c7d3b3
commit 7c41566
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/ibis/backends/tests/test_json.py b/ibis/backends/tests/test_json.py
@@ -5,7 +5,7 @@
 from pytest import param
 
 
-@pytest.mark.notimpl(["bigquery", "datafusion", "pyspark", "mssql", "trino"])
+@pytest.mark.notimpl(["bigquery", "datafusion", "pyspark", "mssql"])
 @pytest.mark.notyet(["clickhouse"], reason="upstream is broken")
 @pytest.mark.never(["impala"], reason="doesn't support JSON and never will")
 @pytest.mark.parametrize(

diff --git a/ibis/backends/trino/registry.py b/ibis/backends/trino/registry.py
@@ -19,6 +19,13 @@ def _arbitrary(t, op):
     return reduction(sa.func.arbitrary)(t, op)
 
 
+def _json_get_item(t, op):
+    arg = t.translate(op.arg)
+    index = t.translate(op.index)
+    fmt = "%d" if op.index.output_dtype.is_integer() else '"%s"'
+    return sa.func.json_extract(arg, sa.func.format(f"$[{fmt}]", index))
+
+
 operation_registry.update(
     {
         # conditional expressions
@@ -36,5 +43,6 @@ def _arbitrary(t, op):
         ops.Covariance: _covar,
         ops.ExtractMillisecond: unary(sa.func.millisecond),
         ops.Arbitrary: _arbitrary,
+        ops.JSONGetItem: _json_get_item,
     }
 )