Fix: Handle single-column DataFrame dtypes in TableWidget

Author: shuoweil

bigframes/display/anywidget.py

@@ -110,10 +110,14 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # TODO(b/469861913): Nested columns from structs (e.g., 'struct_col.name') are not currently sortable.
         # TODO(b/463754889): Support non-string column labels for sorting.
         if all(isinstance(col, str) for col in dataframe.columns):
+            # TODO(b/462525985): dataframe.dtypes should always be a Series.
+            items = (
+                dataframe.dtypes.items()
+                if isinstance(dataframe.dtypes, pd.Series)
+                else [(dataframe.columns[0], dataframe.dtypes)]
+            )
             self.orderable_columns = [
-                str(col_name)
-                for col_name, dtype in dataframe.dtypes.items()
-                if dtypes.is_orderable(dtype)
+                str(col_name) for col_name, dtype in items if dtypes.is_orderable(dtype)
             ]
         else:
             self.orderable_columns = []

tests/unit/display/test_anywidget.py

@@ -153,3 +153,32 @@ def test_page_size_change_resets_sort(mock_df):
 
     # to_pandas_batches called again (reset)
     assert mock_df.to_pandas_batches.call_count >= 2
+
+
+def test_table_widget_single_column_dataframe():
+    """Test that TableWidget can be created with a single-column DataFrame."""
+    from bigframes.display.anywidget import TableWidget
+    import bigframes.dtypes
+
+    mock_df = mock.create_autospec(bigframes.dataframe.DataFrame, instance=True)
+    mock_df.columns = ["image"]
+    # This is the key part: simulate single-column DataFrame's dtypes property
+    mock_df.dtypes = bigframes.dtypes.STRING_DTYPE  # Using string as a stand-in
+
+    mock_block = mock.Mock()
+    mock_block.has_index = False
+    mock_df._block = mock_block
+
+    # We mock _initial_load to avoid complex setup
+    with mock.patch.object(TableWidget, "_initial_load"):
+        with bigframes.option_context(
+            "display.repr_mode", "anywidget", "display.max_rows", 10
+        ):
+            try:
+                widget = TableWidget(mock_df)
+                # Check if 'image' column is considered orderable (it is a string dtype)
+                assert "image" in widget.orderable_columns
+            except IndexError:
+                pytest.fail(
+                    "Creating TableWidget with single-column DataFrame raised IndexError"
+                )