Remove support for pyarrow extension types on the Ray runner.

daft/datatype.py

@@ -4,6 +4,7 @@
 
 import pyarrow as pa
 
+from daft.context import get_context
 from daft.daft import PyDataType
 
 _RAY_DATA_EXTENSIONS_AVAILABLE = True
@@ -178,6 +179,12 @@ def from_arrow_type(cls, arrow_type: pa.lib.DataType) -> DataType:
                 f"used in non-Python Arrow implementations and Daft uses the Rust Arrow2 implementation: {arrow_type}"
             )
         elif isinstance(arrow_type, pa.BaseExtensionType):
+            if get_context().runner_config.name == "ray":
+                raise ValueError(
+                    f"pyarrow extension types are not supported for the Ray runner: {arrow_type}. If you need support "
+                    "for this, please let us know on this issue: "
+                    "https://github.com/Eventual-Inc/Daft/issues/933"
+                )
             name = arrow_type.extension_name
             try:
                 metadata = arrow_type.__arrow_ext_serialize__().decode()

tests/dataframe/test_creation.py

@@ -15,6 +15,7 @@
 
 import daft
 from daft.api_annotations import APITypeError
+from daft.context import get_context
 from daft.dataframe import DataFrame
 from daft.datatype import DataType
 from tests.conftest import UuidType
@@ -173,6 +174,10 @@ def test_create_dataframe_arrow_tensor_ray(valid_data: list[dict[str, float]]) -
     ARROW_VERSION < (12, 0, 0),
     reason=f"Arrow version {ARROW_VERSION} doesn't support the canonical tensor extension type.",
 )
+@pytest.mark.skipif(
+    get_context().runner_config.name == "ray",
+    reason="Pickling canonical tensor extension type is not supported by pyarrow",
+)
 def test_create_dataframe_arrow_tensor_canonical(valid_data: list[dict[str, float]]) -> None:
     pydict = {k: [item[k] for item in valid_data] for k in valid_data[0].keys()}
     dtype = pa.fixed_shape_tensor(pa.int64(), (2, 2))
@@ -191,6 +196,10 @@ def test_create_dataframe_arrow_tensor_canonical(valid_data: list[dict[str, floa
     assert df.to_arrow() == expected
 
 
+@pytest.mark.skipif(
+    get_context().runner_config.name == "ray",
+    reason="pyarrow extension types aren't supported on Ray clusters.",
+)
 def test_create_dataframe_arrow_extension_type(valid_data: list[dict[str, float]], uuid_ext_type: UuidType) -> None:
     pydict = {k: [item[k] for item in valid_data] for k in valid_data[0].keys()}
     storage = pa.array([f"{i}".encode() for i in range(len(valid_data))])
@@ -207,6 +216,22 @@ def test_create_dataframe_arrow_extension_type(valid_data: list[dict[str, float]
     assert df.to_arrow() == expected
 
 
+# TODO(Clark): Remove this test once pyarrow extension types are supported for Ray clusters.
+@pytest.mark.skipif(
+    get_context().runner_config.name != "ray",
+    reason="This test requires the Ray runner.",
+)
+def test_create_dataframe_arrow_extension_type_fails_for_ray(
+    valid_data: list[dict[str, float]], uuid_ext_type: UuidType
+) -> None:
+    pydict = {k: [item[k] for item in valid_data] for k in valid_data[0].keys()}
+    storage = pa.array([f"{i}".encode() for i in range(len(valid_data))])
+    pydict["obj"] = pa.ExtensionArray.from_storage(uuid_ext_type, storage)
+    t = pa.Table.from_pydict(pydict)
+    with pytest.raises(ValueError):
+        daft.from_arrow(t).to_arrow()
+
+
 class PyExtType(pa.PyExtensionType):
     def __init__(self):
         pa.PyExtensionType.__init__(self, pa.binary())