From aae28e9e933a4a26cc928c17563e3647bad06913 Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Wed, 11 Oct 2023 09:35:51 -0400
Subject: [PATCH] feat(duckdb): implement partial json collection casting

---
 ibis/backends/duckdb/registry.py | 10 +++++++---
 ibis/backends/tests/test_json.py |  4 ----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/ibis/backends/duckdb/registry.py b/ibis/backends/duckdb/registry.py
index 1ad137e0113a..332ddbda8250 100644
--- a/ibis/backends/duckdb/registry.py
+++ b/ibis/backends/duckdb/registry.py
@@ -331,6 +331,11 @@ def _try_cast(t, op):
 )
 
 
+def _to_json_collection(t, op):
+    typ = t.get_sqla_type(op.dtype)
+    return try_cast(t.translate(op.arg), typ, type_=typ)
+
+
 operation_registry.update(
     {
         ops.ArrayColumn: (
@@ -477,6 +482,8 @@ def _try_cast(t, op):
         ops.TimeDelta: _temporal_delta,
         ops.DateDelta: _temporal_delta,
         ops.TimestampDelta: _temporal_delta,
+        ops.ToJSONMap: _to_json_collection,
+        ops.ToJSONArray: _to_json_collection,
     }
 )
 
@@ -486,9 +493,6 @@ def _try_cast(t, op):
     ops.NTile,
     # ibis.expr.operations.strings
     ops.Translate,
-    # ibis.expr.operations.json
-    ops.ToJSONMap,
-    ops.ToJSONArray,
 }
 
 operation_registry = {
diff --git a/ibis/backends/tests/test_json.py b/ibis/backends/tests/test_json.py
index 3c50992cdc99..721456966f43 100644
--- a/ibis/backends/tests/test_json.py
+++ b/ibis/backends/tests/test_json.py
@@ -9,8 +9,6 @@
 from packaging.version import parse as vparse
 from pytest import param
 
-from ibis.common.exceptions import OperationNotDefinedError
-
 pytestmark = [
     pytest.mark.never(["impala"], reason="doesn't support JSON and never will"),
     pytest.mark.notyet(["clickhouse"], reason="upstream is broken"),
@@ -50,7 +48,6 @@ def test_json_getitem(json_t, expr_fn, expected):
 @pytest.mark.notyet(
     ["pyspark", "trino"], reason="should work but doesn't deserialize JSON"
 )
-@pytest.mark.notimpl(["duckdb"], raises=OperationNotDefinedError)
 def test_json_map(json_t):
     expr = json_t.js.map.name("res")
     result = expr.execute()
@@ -75,7 +72,6 @@ def test_json_map(json_t):
     ["pyspark", "trino"], reason="should work but doesn't deserialize JSON"
 )
 @pytest.mark.notyet(["bigquery"], reason="doesn't allow null in arrays")
-@pytest.mark.notimpl(["duckdb"], raises=OperationNotDefinedError)
 def test_json_array(json_t):
     expr = json_t.js.array.name("res")
     result = expr.execute()