Repository: ibis-project/ibis

Issue:
bug(postgres): `to_pyarrow` fails with json type
### What happened?

Hi there -- I'm trying to convert a postgres table to pyarrow and getting this error:

`ArrowNotImplementedError: extension`

I went into the ibis backend to look at what Ibis thinks the types are, and found this is what it thinks the array types are

```
struct<id: string not null, created_at: timestamp[us, tz=UTC], response: extension<ibis.json<JSONType>>, email: string>
```

And here is the ibis expression:

```
r0 := DatabaseTable: waitlist
  id         !uuid
  created_at timestamp('UTC')
  response   json
  email      string
Limit[r0, n=10000]
```

### What version of ibis are you using?

8.0.0

### What backend(s) are you using, if any?

Postgres

### Relevant log output

```sh
--------------------------------------------------------------------------
ArrowNotImplementedError                  Traceback (most recent call last)
Cell In[20], line 1
----> 1 w.execution_helper().to_pyarrow()

File ~/Documents/GitHub/spoonbill/ibis/expr/types/core.py:444, in Expr.to_pyarrow(self, params, limit, **kwargs)
    416 @experimental
    417 def to_pyarrow(
    418     self,
   (...)
    422     **kwargs: Any,
    423 ) -> pa.Table:
    424     """Execute expression and return results in as a pyarrow table.
    425
    426     This method is eager and will execute the associated expression
   (...)
    442         A pyarrow table holding the results of the executed expression.
    443     """
--> 444     return self._find_backend(use_default=True).to_pyarrow(
    445         self, params=params, limit=limit, **kwargs
    446     )

File ~/Documents/GitHub/spoonbill/ibis/backends/base/__init__.py:367, in _FileIOHandler.to_pyarrow(self, expr, params, limit, **kwargs)
    363 arrow_schema = schema.to_pyarrow()
    364 with self.to_pyarrow_batches(
    365     table_expr, params=params, limit=limit, **kwargs
    366 ) as reader:
--> 367     table = pa.Table.from_batches(reader, schema=arrow_schema)
    369 return expr.__pyarrow_result__(
    370     table.rename_columns(table_expr.columns).cast(arrow_schema)
    371 )

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/table.pxi:4104, in pyarrow.lib.Table.from_batches()

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/ipc.pxi:666, in pyarrow.lib.RecordBatchReader.__next__()

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/ipc.pxi:700, in pyarrow.lib.RecordBatchReader.read_next_batch()

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/types.pxi:88, in pyarrow.lib._datatype_to_pep3118()

File ~/Documents/GitHub/spoonbill/ibis/backends/base/sql/__init__.py:245, in <genexpr>(.0)
    242 array_type = schema.as_struct().to_pyarrow()
    243 print(array_type)
    244 arrays = (
--> 245     pa.array(map(tuple, batch), type=array_type)
    246     for batch in self._cursor_batches(
    247         expr, params=params, limit=limit, chunk_size=chunk_size
    248     )
    249 )
    250 batches = map(pa.RecordBatch.from_struct_array, arrays)
    252 return pa.ipc.RecordBatchReader.from_batches(schema.to_pyarrow(), batches)

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/array.pxi:344, in pyarrow.lib.array()

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/array.pxi:42, in pyarrow.lib._sequence_to_array()

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/error.pxi:154, in pyarrow.lib.pyarrow_internal_check_status()

File ~/Library/Caches/pypoetry/virtualenvs/vinyl-cwr2Pa_2-py3.11/lib/python3.11/site-packages/pyarrow/error.pxi:91, in pyarrow.lib.check_status()
```


### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
@cpcloud -- confirmed that this still exists with v. 9.0.0
Yep, this is a problem in the postgres `to_pyarrow` implementation, which wasn't touched much in the latest big refactor.
I think the underlying execute() function also uses to_pyarrow, so effectively the connector is blocked. Is there an alternative that will work? What is the ETA for fixing this?
You should be able to use something like

```python
pa.Table.from_pandas(expr.to_pandas())
```

If you definitely need a PyArrow Table.
Not 100% sure what the ETA is. It will probably be in the next release, but no promises :)
Ok will try that — thank you!

> On Feb 21, 2024, at 11:02 AM, Phillip Cloud ***@***.***> wrote:
> 
> 
> Not 100% sure what the ETA is. It will probably be in the next release, but no promises :)
> 
> —
> Reply to this email directly, view it on GitHub <https://github.com/ibis-project/ibis/issues/8318#issuecomment-1957315636>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/A5HSLR5UIOPIULYAOJDQVLLYUYSBLAVCNFSM6AAAAABDFFFD36VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNJXGMYTKNRTGY>.
> You are receiving this because you authored the thread.
> 



Resolving patch:
diff --git a/ibis/backends/postgres/__init__.py b/ibis/backends/postgres/__init__.py
--- a/ibis/backends/postgres/__init__.py
+++ b/ibis/backends/postgres/__init__.py
@@ -260,7 +260,7 @@ def do_connect(
             month : int32
 
         """
-
+        psycopg2.extras.register_default_json(loads=lambda x: x)
         self.con = psycopg2.connect(
             host=host,
             port=port,
diff --git a/ibis/backends/snowflake/converter.py b/ibis/backends/snowflake/converter.py
--- a/ibis/backends/snowflake/converter.py
+++ b/ibis/backends/snowflake/converter.py
@@ -1,18 +1,54 @@
 from __future__ import annotations
 
 import datetime
+import json
 from typing import TYPE_CHECKING
 
+import pyarrow as pa
+
 from ibis.formats.pandas import PandasData
-from ibis.formats.pyarrow import PYARROW_JSON_TYPE, PyArrowData
+from ibis.formats.pyarrow import PyArrowData
 
 if TYPE_CHECKING:
-    import pyarrow as pa
-
     import ibis.expr.datatypes as dt
     from ibis.expr.schema import Schema
 
 
+class JSONScalar(pa.ExtensionScalar):
+    def as_py(self):
+        value = self.value
+        if value is None:
+            return value
+        else:
+            return json.loads(value.as_py())
+
+
+class JSONArray(pa.ExtensionArray):
+    pass
+
+
+class JSONType(pa.ExtensionType):
+    def __init__(self):
+        super().__init__(pa.string(), "ibis.json")
+
+    def __arrow_ext_serialize__(self):
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(cls, storage_type, serialized):
+        return cls()
+
+    def __arrow_ext_class__(self):
+        return JSONArray
+
+    def __arrow_ext_scalar_class__(self):
+        return JSONScalar
+
+
+PYARROW_JSON_TYPE = JSONType()
+pa.register_extension_type(PYARROW_JSON_TYPE)
+
+
 class SnowflakePandasData(PandasData):
     @classmethod
     def convert_Timestamp_element(cls, dtype):
@@ -50,16 +86,12 @@ def convert_Struct(cls, s, dtype, pandas_type):
 class SnowflakePyArrowData(PyArrowData):
     @classmethod
     def convert_table(cls, table: pa.Table, schema: Schema) -> pa.Table:
-        import pyarrow as pa
-
         columns = [cls.convert_column(table[name], typ) for name, typ in schema.items()]
         return pa.Table.from_arrays(columns, names=schema.names)
 
     @classmethod
     def convert_column(cls, column: pa.Array, dtype: dt.DataType) -> pa.Array:
         if dtype.is_json() or dtype.is_array() or dtype.is_map() or dtype.is_struct():
-            import pyarrow as pa
-
             if isinstance(column, pa.ChunkedArray):
                 column = column.combine_chunks()
 
diff --git a/ibis/formats/pyarrow.py b/ibis/formats/pyarrow.py
--- a/ibis/formats/pyarrow.py
+++ b/ibis/formats/pyarrow.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 from typing import TYPE_CHECKING, Any
 
 import pyarrow as pa
@@ -14,37 +13,6 @@
     from collections.abc import Sequence
 
 
-class JSONScalar(pa.ExtensionScalar):
-    def as_py(self):
-        value = self.value
-        if value is None:
-            return value
-        else:
-            return json.loads(value.as_py())
-
-
-class JSONArray(pa.ExtensionArray):
-    pass
-
-
-class JSONType(pa.ExtensionType):
-    def __init__(self):
-        super().__init__(pa.string(), "ibis.json")
-
-    def __arrow_ext_serialize__(self):
-        return b""
-
-    @classmethod
-    def __arrow_ext_deserialize__(cls, storage_type, serialized):
-        return cls()
-
-    def __arrow_ext_class__(self):
-        return JSONArray
-
-    def __arrow_ext_scalar_class__(self):
-        return JSONScalar
-
-
 _from_pyarrow_types = {
     pa.int8(): dt.Int8,
     pa.int16(): dt.Int16,
@@ -91,6 +59,7 @@ def __arrow_ext_scalar_class__(self):
     dt.MACADDR: pa.string(),
     dt.INET: pa.string(),
     dt.UUID: pa.string(),
+    dt.JSON: pa.string(),
 }
 
 
@@ -128,8 +97,6 @@ def to_ibis(cls, typ: pa.DataType, nullable=True) -> dt.DataType:
             key_dtype = cls.to_ibis(typ.key_type, typ.key_field.nullable)
             value_dtype = cls.to_ibis(typ.item_type, typ.item_field.nullable)
             return dt.Map(key_dtype, value_dtype, nullable=nullable)
-        elif isinstance(typ, JSONType):
-            return dt.JSON()
         elif pa.types.is_dictionary(typ):
             return cls.to_ibis(typ.value_type)
         else:
@@ -191,8 +158,6 @@ def from_ibis(cls, dtype: dt.DataType) -> pa.DataType:
                 nullable=dtype.value_type.nullable,
             )
             return pa.map_(key_field, value_field, keys_sorted=False)
-        elif dtype.is_json():
-            return PYARROW_JSON_TYPE
         elif dtype.is_geospatial():
             return pa.binary()
         else:
@@ -303,7 +268,3 @@ def to_frame(self):
 
     def to_pyarrow(self, schema: Schema) -> pa.Table:
         return self.obj
-
-
-PYARROW_JSON_TYPE = JSONType()
-pa.register_extension_type(PYARROW_JSON_TYPE)


Test patch:
diff --git a/ibis/backends/snowflake/tests/conftest.py b/ibis/backends/snowflake/tests/conftest.py
--- a/ibis/backends/snowflake/tests/conftest.py
+++ b/ibis/backends/snowflake/tests/conftest.py
@@ -119,27 +119,29 @@ def add_catalog_and_schema(node):
     def _load_data(self, **_: Any) -> None:
         """Load test data into a Snowflake backend instance."""
 
+        connect_args = {}
+
         url = urlparse(_get_url())
-        db, schema = url.path[1:].split("/", 1)
-        (warehouse,) = parse_qs(url.query)["warehouse"]
-        connect_args = {
-            "user": url.username,
-            "password": url.password,
-            "account": url.hostname,
-            "warehouse": warehouse,
-        }
-
-        session_parameters = {
-            "MULTI_STATEMENT_COUNT": 0,
-            "JSON_INDENT": 0,
-            "PYTHON_CONNECTOR_QUERY_RESULT_FORMAT": "arrow_force",
-        }
+
+        if url.path:
+            db, schema = url.path[1:].split("/", 1)
+            (warehouse,) = parse_qs(url.query)["warehouse"]
+            connect_args.update(
+                {
+                    "user": url.username,
+                    "password": url.password,
+                    "account": url.hostname,
+                    "warehouse": warehouse,
+                }
+            )
+        else:
+            db = os.environ["SNOWFLAKE_DATABASE"]
+            schema = os.environ["SNOWFLAKE_SCHEMA"]
 
         dbschema = f"{db}.{schema}"
 
-        with closing(
-            sc.connect(**connect_args, session_parameters=session_parameters)
-        ) as con, closing(con.cursor()) as c:
+        with closing(sc.connect(**connect_args)) as con, closing(con.cursor()) as c:
+            c.execute("ALTER SESSION SET MULTI_STATEMENT_COUNT = 0 JSON_INDENT = 0")
             c.execute(
                 f"""
                 CREATE DATABASE IF NOT EXISTS {db};
diff --git a/ibis/backends/snowflake/tests/test_client.py b/ibis/backends/snowflake/tests/test_client.py
--- a/ibis/backends/snowflake/tests/test_client.py
+++ b/ibis/backends/snowflake/tests/test_client.py
@@ -304,3 +304,16 @@ def test_no_argument_connection():
 
     con = ibis.connect("snowflake://")
     assert con.list_tables() is not None
+
+
+def test_struct_of_json(con):
+    raw = {"a": [1, 2, 3], "b": "456"}
+    lit = ibis.struct(raw)
+    expr = lit.cast("struct<a: array<int>, b: json>")
+
+    n = 5
+    t = con.tables.functional_alltypes.mutate(lit=expr).limit(n).lit
+    result = con.to_pyarrow(t)
+
+    assert len(result) == n
+    assert all(value == raw for value in result.to_pylist())
diff --git a/ibis/backends/tests/test_client.py b/ibis/backends/tests/test_client.py
--- a/ibis/backends/tests/test_client.py
+++ b/ibis/backends/tests/test_client.py
@@ -4,6 +4,7 @@
 import contextlib
 import importlib
 import inspect
+import json
 import re
 import string
 import subprocess
@@ -1468,3 +1469,46 @@ def test_close_connection(con):
     # DB-API states that subsequent execution attempt should raise
     with pytest.raises(Exception):  # noqa:B017
         new_con.list_tables()
+
+
+@pytest.mark.notyet(
+    ["clickhouse"],
+    raises=AttributeError,
+    reason="JSON extension is experimental and not enabled by default in testing",
+)
+@pytest.mark.notyet(
+    ["datafusion", "polars", "mssql", "druid", "oracle", "exasol", "impala"],
+    raises=AttributeError,
+    reason="JSON type not implemented",
+)
+@pytest.mark.notimpl(
+    ["risingwave", "sqlite"],
+    raises=pa.ArrowTypeError,
+    reason="mismatch between output value and expected input type",
+)
+@pytest.mark.never(
+    ["snowflake"],
+    raises=TypeError,
+    reason="snowflake uses a custom pyarrow extension type for JSON pretty printing",
+)
+def test_json_to_pyarrow(con):
+    t = con.tables.json_t
+    table = t.to_pyarrow()
+    js = table["js"]
+
+    expected = [
+        {"a": [1, 2, 3, 4], "b": 1},
+        {"a": None, "b": 2},
+        {"a": "foo", "c": None},
+        None,
+        [42, 47, 55],
+        [],
+    ]
+    expected = {json.dumps(val) for val in expected}
+
+    result = {
+        # loads and dumps so the string representation is the same
+        json.dumps(json.loads(val))
+        for val in js.to_pylist()
+    }
+    assert result == expected