disable sorting for integer or multiindex

Author: shuoweil

bigframes/display/anywidget.py

@@ -108,11 +108,15 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # set traitlets properties that trigger observers
         # TODO(b/462525985): Investigate and improve TableWidget UX for DataFrames with a large number of columns.
         self.page_size = initial_page_size
-        self.orderable_columns = [
-            str(col_name)
-            for col_name, dtype in dataframe.dtypes.items()
-            if dtypes.is_orderable(dtype)
-        ]
+        # TODO(b/463754889): Support non-string column labels for sorting.
+        if all(isinstance(col, str) for col in dataframe.columns):
+            self.orderable_columns = [
+                col_name
+                for col_name, dtype in dataframe.dtypes.items()
+                if dtypes.is_orderable(dtype)
+            ]
+        else:
+            self.orderable_columns = []
 
         # obtain the row counts
         # TODO(b/428238610): Start iterating over the result of `to_pandas_batches()`
@@ -262,6 +266,7 @@ def _set_table_html(self) -> None:
         # Apply sorting if a column is selected
         df_to_display = self._dataframe
         if self.sort_column:
+            # TODO(b/463715504): Support sorting by index columns.
             df_to_display = df_to_display.sort_values(
                 by=self.sort_column, ascending=self.sort_ascending
             )

notebooks/dataframes/anywidget_mode.ipynb

@@ -157,7 +157,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f0f1cc04e1564a6ab4ba39935aebd11a",
+       "model_id": "3ad3ce907f4d43de8c3c15f0d1cdebd6",
        "version_major": 2,
        "version_minor": 1
       },
@@ -255,7 +255,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bed78e341154486bb49841ce12a0af93",
+       "model_id": "684becdc45de4a15826672c0a7e89930",
        "version_major": 2,
        "version_minor": 1
       },
@@ -369,7 +369,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "357c732c110e44fa9efe3e512da61638",
+       "model_id": "5cccf9eb2a674f9f8a1bdb165921d5f9",
        "version_major": 2,
        "version_minor": 1
       },
@@ -409,7 +409,7 @@
      "data": {
       "text/html": [
        "✅ Completed. \n",
-       "    Query processed 85.9 kB in 14 seconds of slot time.\n",
+       "    Query processed 85.9 kB in 13 seconds of slot time.\n",
        "    "
       ],
       "text/plain": [
@@ -456,7 +456,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "af87c96ce359412a99b5eabbc1b9d137",
+       "model_id": "06e704906f254de194e5e3ea1ef2bdbb",
        "version_major": 2,
        "version_minor": 1
       },

tests/system/small/test_anywidget.py

@@ -721,11 +721,6 @@ def test_widget_with_unknown_row_count_empty_dataframe(
         assert widget.page == 0
 
 
-# TODO(shuowei): Add tests for custom index and multiindex
-# This may not be necessary for the SQL Cell use case but should be
-# considered for completeness.
-
-
 def test_widget_sort_should_sort_ascending_on_first_click(
     table_widget, paginated_pandas_df: pd.DataFrame
 ):
@@ -817,3 +812,55 @@ def test_widget_sort_should_reset_on_page_size_change(
     html = table_widget.table_html
 
     _assert_html_matches_pandas_slice(html, expected_slice, paginated_pandas_df)
+
+
+@pytest.fixture(scope="module")
+def integer_column_df(session):
+    """Create a DataFrame with integer column labels."""
+    pandas_df = pd.DataFrame([[0, 1], [2, 3]], columns=pd.Index([1, 2]))
+    return session.read_pandas(pandas_df)
+
+
+@pytest.fixture(scope="module")
+def multiindex_column_df(session):
+    """Create a DataFrame with MultiIndex column labels."""
+    pandas_df = pd.DataFrame(
+        {
+            "foo": ["one", "one", "one", "two", "two", "two"],
+            "bar": ["A", "B", "C", "A", "B", "C"],
+            "baz": [1, 2, 3, 4, 5, 6],
+            "zoo": ["x", "y", "z", "q", "w", "t"],
+        }
+    )
+    df = session.read_pandas(pandas_df)
+    # The session is attached to `df` through the constructor.
+    # We can pass it to the pivoted DataFrame.
+    pdf = df.pivot(index="foo", columns="bar", values=["baz", "zoo"])
+    return pdf
+
+
+def test_table_widget_integer_columns_disables_sorting(integer_column_df):
+    """
+    Given a DataFrame with integer column labels, the widget should
+    disable sorting.
+    """
+    from bigframes.display import TableWidget
+
+    widget = TableWidget(integer_column_df)
+    assert widget.orderable_columns == []
+
+
+def test_table_widget_multiindex_columns_disables_sorting(multiindex_column_df):
+    """
+    Given a DataFrame with a MultiIndex for columns, the widget should
+    disable sorting.
+    """
+    from bigframes.display import TableWidget
+
+    widget = TableWidget(multiindex_column_df)
+    assert widget.orderable_columns == []
+
+
+# TODO(shuowei): Add tests for custom index and multiindex
+# This may not be necessary for the SQL Cell use case but should be
+# considered for completeness.