fix(postgres): support uuid.to_pyarrow()

Author: NickCrews

ibis/backends/postgres/__init__.py

@@ -26,6 +26,7 @@
 from ibis.backends import CanCreateDatabase, CanListCatalog, PyArrowExampleLoader
 from ibis.backends.sql import SQLBackend
 from ibis.backends.sql.compilers.base import TRUE, C, ColGen
+from ibis.formats.pyarrow import to_pa_compatible
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping
@@ -808,6 +809,7 @@ def _batches(self: Self, *, schema: pa.Schema, query: str):
 
         self._run_pre_execute_hooks(expr)
 
+        expr = to_pa_compatible(expr)
         schema = expr.as_table().schema().to_pyarrow()
         query = self.compile(expr, limit=limit, params=params)
         return pa.RecordBatchReader.from_batches(

ibis/backends/tests/test_uuid.py

@@ -67,3 +67,10 @@ def test_uuid_unique_each_row(con):
         con.tables.functional_alltypes.mutate(uuid=ibis.uuid()).limit(2).uuid.nunique()
     )
     assert expr.execute() == 2
+
+
+def test_uuid_pyarrow(con):
+    pa = pytest.importorskip("pyarrow")
+    expr = ibis.literal(RAW_TEST_UUID, type=dt.uuid)
+    result = con.to_pyarrow(expr)
+    assert result == pa.scalar(RAW_TEST_UUID)

ibis/formats/pyarrow.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import pyarrow as pa
 import pyarrow_hotfix  # noqa: F401
@@ -19,6 +19,10 @@
     import polars as pl
     import pyarrow.dataset as ds
 
+    import ibis
+
+    TableOrValue = TypeVar("TableOrValue", bound=ibis.Table | ibis.Value)
+
 
 _from_pyarrow_types = {
     pa.int8(): dt.Int8,
@@ -384,3 +388,50 @@ def to_pyarrow_dataset(self, schema: Schema) -> ds.Dataset:
 
     def to_polars(self, schema: Schema) -> pa.Table:
         raise com.UnsupportedOperationError(self.ERROR_MESSAGE)
+
+
+def to_pa_compatible(table_or_val: TableOrValue) -> TableOrValue:
+    """Convert (on the backend) an Ibis table or value to a PyArrow compatible type.
+
+    If we have a uuid type on the backend, we are going to represent it on the
+    pyarrow side as a string. So, since we are going to cast it to a string anyway,
+    we might as well do it on the backend side. This is a performance gain.
+
+    Not only that, but it also avoids some issues where we have a type (eg uuid)
+    which is not supported by pyarrow, and we run into trouble materializing
+    it to pyarrow. See https://github.com/ibis-project/ibis/issues/8532
+    """
+    import ibis
+
+    if isinstance(table_or_val, ibis.Table):
+        return _to_pa_compatible_table(table_or_val)
+    elif isinstance(table_or_val, ibis.Value):
+        return _to_pa_compatible_value(table_or_val)
+    else:
+        raise TypeError(f"Unsupported type: {type(table_or_val)}")
+
+
+def _to_pa_compatible_value(val: ibis.Value) -> ibis.Value:
+    original_type = val.type()
+    # Convert the original type back and forth to check biyectivity
+    # of types mappings pyarrow <-> ibis.
+    pa_compatible_type = PyArrowType.to_ibis(PyArrowType.from_ibis(original_type))
+    if original_type != pa_compatible_type:
+        # If the original type is not compatible with PyArrow, we cast
+        # server side to match the types that PyArrow expects in Ibis.
+        val = val.cast(pa_compatible_type)
+    return val
+
+
+def _to_pa_compatible_table(table: ibis.Table) -> ibis.Table:
+    original_schema = table.schema()
+    # Convert the original schema back and forth to check biyectivity
+    # of types mappings pyarrow <-> ibis.
+    pa_compatible_schema = PyArrowSchema.to_ibis(
+        PyArrowSchema.from_ibis(original_schema)
+    )
+    if original_schema != pa_compatible_schema:
+        # If the original schema is not compatible with PyArrow, we cast
+        # server side to match the types that PyArrow expects in Ibis.
+        table = table.cast(pa_compatible_schema)
+    return table