cast list of enums to string preemptively for polars

marimo/_plugins/ui/_impl/tables/polars_table.py

@@ -118,6 +118,15 @@ def to_json_str(
                             result = self._convert_time_to_string(
                                 result, column
                             )
+                        # https://github.com/marimo-team/marimo/issues/7032
+                        # Polars issue with ordering and write_json for enums, so we convert to strings
+                        elif isinstance(dtype, pl.List) and isinstance(
+                            dtype.inner, (pl.Enum, pl.Categorical)
+                        ):
+                            # Convert each element in the list to a string
+                            result = result.with_columns(
+                                pl.col(column.name).cast(pl.List(pl.String))
+                            )
                     return sanitize_json_bigint(result.write_json())
                 except (
                     BaseException
@@ -154,16 +163,6 @@ def to_json_str(
                                 result, column
                             )
                             converted_columns.append(column.name)
-                        # https://github.com/marimo-team/marimo/issues/5562
-                        elif isinstance(dtype, pl.List) and isinstance(
-                            dtype.inner, (pl.Enum, pl.Categorical)
-                        ):
-                            # Convert each element in the list to a string
-                            result = result.with_columns(
-                                pl.col(column.name).cast(pl.List(pl.String))
-                            )
-                            converted_columns.append(column.name)
-
                     if converted_columns:
                         LOGGER.info(
                             "Converted columns %s to safe values.",

tests/_plugins/ui/_impl/tables/test_polars_table.py

@@ -3,6 +3,7 @@
 import datetime
 import json
 import unittest
+from enum import Enum
 from math import isnan
 from typing import Any
 
@@ -995,3 +996,35 @@ def test_to_json_enum_list_supported(self) -> None:
 
         data_list = pl.DataFrame(data, schema={"A": pl.List(pl.Categorical())})
         data_list.write_json()
+
+    @pytest.mark.xfail(
+        reason="Polars does not properly order sliced data to json when enums in list"
+    )
+    def test_failing_enums_list(self) -> None:
+        import polars as pl
+
+        class MyEnum(Enum):
+            A = 1
+            B = 2
+            C = 3
+            D = 4
+
+        # Create 10 rows cycling through enum values B, C, D, A...
+        enum_names = [e.name for e in MyEnum]
+        rows = [
+            {"value": [enum_names[i % len(enum_names)]]} for i in range(1, 11)
+        ]
+
+        expected_first_five = '[{"value":["B"]},{"value":["C"]},{"value":["D"]},{"value":["A"]},{"value":["B"]}]'
+        expected_second_five = '[{"value":["C"]},{"value":["D"]},{"value":["A"]},{"value":["B"]},{"value":["C"]}]'
+
+        # Test without schema - works fine
+        df = pl.DataFrame(rows)
+        assert df[0:5].write_json() == expected_first_five
+        assert df[5:10].write_json() == expected_second_five
+
+        # Test with schema - second slice fails
+        schema = {"value": pl.List(pl.Enum(enum_names))}
+        df_schema = pl.DataFrame(rows, schema=schema)
+        assert df_schema[0:5].write_json() == expected_first_five
+        assert df_schema[5:10].write_json() == expected_second_five  # fails

tests/_plugins/ui/_impl/test_table.py

@@ -3,6 +3,7 @@
 
 import json
 from datetime import date
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 from unittest.mock import patch
 
@@ -2399,3 +2400,41 @@ def hover_text(_row: str, _col: str, value: Any) -> str:
         "4": {"column_0": "hover:carrots"},
         "5": {"column_0": "hover:carrots"},
     }
+
+
+@pytest.mark.skipif(
+    not DependencyManager.polars.has(), reason="Polars not installed"
+)
+def test_polars_enums_in_list():
+    import polars as pl
+
+    class MyEnum(Enum):
+        A = 1
+        B = 2
+        C = 3
+        D = 4
+
+    # Create 10 rows cycling through enum values B, C, D, A...
+    enum_names = [e.name for e in MyEnum]
+    rows = [{"value": [enum_names[i % len(enum_names)]]} for i in range(1, 11)]
+
+    schema = {"value": pl.List(pl.Enum(enum_names))}
+    df = pl.DataFrame(rows, schema=schema)
+
+    table = ui.table(df, selection=None)
+
+    # First page
+    response = table._search(SearchTableArgs(page_size=5, page_number=0))
+    assert (
+        response.data
+        == '[{"value":["B"]},{"value":["C"]},{"value":["D"]},{"value":["A"]},{"value":["B"]}]'
+    )
+
+    # Second page
+    response_next_page = table._search(
+        SearchTableArgs(page_size=5, page_number=1)
+    )
+    assert (
+        response_next_page.data
+        == '[{"value":["C"]},{"value":["D"]},{"value":["A"]},{"value":["B"]},{"value":["C"]}]'
+    )