fix(polars): support polars Enum type (#10017)

ibis/backends/polars/tests/test_client.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import polars as pl
+import polars.testing
 import pytest
 
 import ibis
@@ -37,3 +39,24 @@ def test_array_flatten(con):
         {"id": data["id"], "flat": [row[0] for row in data["happy"]]}
     )
     tm.assert_frame_equal(result.to_pandas(), expected)
+
+
+def test_memtable_polars_types(con):
+    # Check that we can create a memtable with some polars-specific types,
+    # and that those columns then work in downstream operations
+    df = pl.DataFrame(
+        {
+            "x": ["a", "b", "a"],
+            "y": ["c", "d", "c"],
+            "z": ["e", "f", "e"],
+        },
+        schema={
+            "x": pl.String,
+            "y": pl.Categorical,
+            "z": pl.Enum(["e", "f"]),
+        },
+    )
+    t = ibis.memtable(df)
+    res = con.to_polars((t.x + t.y + t.z).name("test"))
+    sol = (df["x"] + df["y"] + df["z"]).rename("test")
+    pl.testing.assert_series_equal(res, sol)

ibis/formats/polars.py

@@ -43,7 +43,7 @@ def to_ibis(cls, typ: pl.DataType, nullable=True) -> dt.DataType:
         """Convert a polars type to an ibis type."""
 
         base_type = typ.base_type()
-        if base_type is pl.Categorical:
+        if base_type in (pl.Categorical, pl.Enum):
             return dt.String(nullable=nullable)
         elif base_type is pl.Decimal:
             return dt.Decimal(

ibis/formats/tests/test_polars.py

@@ -100,8 +100,9 @@ def test_decimal():
     )
 
 
-def test_categorical():
+def test_enum_categorical():
     assert PolarsType.to_ibis(pl.Categorical()) == dt.string
+    assert PolarsType.to_ibis(pl.Enum(["a", "b"])) == dt.string
 
 
 def test_interval_unsupported_unit():