fix(table): allow freezing pandas index columns (#9631)

Author: kirangadhave

marimo/_plugins/ui/_impl/table.py

@@ -756,6 +756,7 @@ def __init__(
         search_result_raw_data: str | None = None
         field_types: FieldTypes | None = None
         num_columns = 0
+        row_headers = self._manager.get_row_headers()
 
         if not _internal_lazy:
             # Search first page
@@ -776,8 +777,12 @@ def __init__(
             # Validate column configurations
             column_names_set = set(self._manager.get_column_names())
             num_columns = len(column_names_set)
+            row_header_names_set = {name for name, _ in row_headers}
             _validate_frozen_columns(
-                freeze_columns_left, freeze_columns_right, column_names_set
+                freeze_columns_left,
+                freeze_columns_right,
+                column_names_set,
+                row_header_names_set,
             )
             _validate_column_formatting(
                 text_justify_columns, wrapped_columns, column_names_set
@@ -816,7 +821,7 @@ def __init__(
                 "show-page-size-selector": show_page_size_selector,
                 "show-column-explorer": show_column_explorer,
                 "show-chart-builder": show_chart_builder,
-                "row-headers": self._manager.get_row_headers(),
+                "row-headers": row_headers,
                 "freeze-columns-left": freeze_columns_left,
                 "freeze-columns-right": freeze_columns_right,
                 "text-justify-columns": text_justify_columns,
@@ -1748,12 +1753,16 @@ def _validate_frozen_columns(
     freeze_columns_left: Sequence[str] | None,
     freeze_columns_right: Sequence[str] | None,
     column_names_set: set[str],
+    row_header_names_set: set[str],
 ) -> None:
     """Validate frozen column configurations.
 
     Validates that:
     1. The same column is not frozen on both sides
-    2. All frozen columns exist in the table
+    2. All left-frozen columns exist as table columns or row-header names
+    3. Right-frozen columns exist as table columns; row-header names are
+       rejected with a friendly error since row headers always render on
+       the left
     """
 
     freeze_columns_left_set = (
@@ -1763,20 +1772,37 @@ def _validate_frozen_columns(
         set(freeze_columns_right) if freeze_columns_right else None
     )
 
-    # Convert sequences to sets for O(1) lookups
     if freeze_columns_left_set and freeze_columns_right_set:
         if not freeze_columns_left_set.isdisjoint(freeze_columns_right_set):
             raise ValueError("The same column cannot be frozen on both sides.")
 
-    # Check all frozen columns exist
     if freeze_columns_left_set:
-        invalid = freeze_columns_left_set - column_names_set
+        # Unnamed row headers (e.g. a default pandas index) have no stable
+        # client-side id, so we can't freeze them. Surface this directly
+        # rather than letting the frontend silently no-op.
+        if "" in freeze_columns_left_set and "" in row_header_names_set:
+            raise ValueError(
+                "Cannot freeze an unnamed row index. "
+                "Set `df.index.name = '...'` (or `df.index.names = [...]` "
+                "for a MultiIndex) and pass that name to freeze_columns_left."
+            )
+        invalid = (
+            freeze_columns_left_set - column_names_set - row_header_names_set
+        )
         if invalid:
             raise ValueError(
                 f"Column '{next(iter(invalid))}' not found in table."
             )
 
     if freeze_columns_right_set:
+        row_header_on_right = freeze_columns_right_set & row_header_names_set
+        if row_header_on_right:
+            name = next(iter(row_header_on_right))
+            raise ValueError(
+                f"Row index '{name}' cannot be frozen on the right; "
+                "row headers always render on the left. "
+                "Use freeze_columns_left instead."
+            )
         invalid = freeze_columns_right_set - column_names_set
         if invalid:
             raise ValueError(

tests/_plugins/ui/_impl/test_table.py

@@ -1112,6 +1112,94 @@ def test_table_with_frozen_columns() -> None:
     assert table._component_args["freeze-columns-right"] == ["d", "e"]
 
 
+@pytest.mark.skipif(
+    not DependencyManager.pandas.has(), reason="Pandas not installed"
+)
+class TestFrozenRowHeaders:
+    def test_freeze_unnamed_pandas_index_rejected(self) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame({"a": [1, 2, 3]}, index=["x", "y", "z"])
+        with pytest.raises(ValueError, match="unnamed row index"):
+            ui.table(df, freeze_columns_left=[""])
+
+    def test_freeze_named_pandas_index(self) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame(
+            {"a": [1, 2]}, index=pd.Index(["x", "y"], name="foo")
+        )
+        table = ui.table(df, freeze_columns_left=["foo"])
+        assert table._component_args["freeze-columns-left"] == ["foo"]
+
+    def test_freeze_multiindex_levels(self) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame(
+            {"v": [1, 2, 3, 4]},
+            index=pd.MultiIndex.from_tuples(
+                [("a", 1), ("a", 2), ("b", 1), ("b", 2)], names=["g", "n"]
+            ),
+        )
+        table = ui.table(df, freeze_columns_left=["g", "n"])
+        assert table._component_args["freeze-columns-left"] == ["g", "n"]
+
+    def test_freeze_collision_suffixed_index(self) -> None:
+        import pandas as pd
+
+        # Index name 'a' collides with a column named 'a'; the row-header
+        # name is suffixed to '_index' (see _resolve_index_name).
+        df = pd.DataFrame({"a": [1, 2]}, index=pd.Index(["x", "y"], name="a"))
+        table = ui.table(df, freeze_columns_left=["a_index"])
+        assert table._component_args["freeze-columns-left"] == ["a_index"]
+
+    def test_freeze_index_and_column_mixed(self) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame(
+            {"a": [1, 2], "b": [3, 4]},
+            index=pd.Index(["x", "y"], name="foo"),
+        )
+        table = ui.table(
+            df,
+            freeze_columns_left=["foo", "a"],
+            freeze_columns_right=["b"],
+        )
+        assert table._component_args["freeze-columns-left"] == ["foo", "a"]
+        assert table._component_args["freeze-columns-right"] == ["b"]
+
+    def test_freeze_row_header_on_right_raises(self) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame(
+            {"a": [1, 2]}, index=pd.Index(["x", "y"], name="foo")
+        )
+        with pytest.raises(
+            ValueError, match="row headers always render on the left"
+        ):
+            ui.table(df, freeze_columns_right=["foo"])
+
+    def test_freeze_unknown_column_still_raises(self) -> None:
+        import pandas as pd
+
+        df = pd.DataFrame(
+            {"a": [1, 2]}, index=pd.Index(["x", "y"], name="foo")
+        )
+        with pytest.raises(ValueError, match="not found in table"):
+            ui.table(df, freeze_columns_left=["nonexistent"])
+
+
+@pytest.mark.skipif(
+    not DependencyManager.polars.has(), reason="Polars not installed"
+)
+def test_freeze_columns_polars_regression() -> None:
+    import polars as pl
+
+    df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    table = ui.table(df, freeze_columns_left=["a"])
+    assert table._component_args["freeze-columns-left"] == ["a"]
+
+
 @pytest.mark.parametrize(
     "df",
     create_dataframes({"a": [1, 2, 3], "b": ["abc", "def", None]}),