feat(display): support multi-column sorting in anywidget

Updates the TableWidget to support sorting by multiple columns.
- Python: Changes sort_column/sort_ascending traits to lists.
- JS: Adds Shift+Click support for multi-column sorting.
- JS: Displays sort indicators for all sorted columns.
- Tests: Adds unit tests for multi-column sorting logic.

Author: shuoweil

bigframes/display/anywidget.py

@@ -53,8 +53,8 @@
 
 @dataclasses.dataclass(frozen=True)
 class _SortState:
-    column: str
-    ascending: bool
+    columns: tuple[str, ...]
+    ascending: tuple[bool, ...]
 
 
 class TableWidget(_WIDGET_BASE):
@@ -68,8 +68,8 @@ class TableWidget(_WIDGET_BASE):
     page_size = traitlets.Int(0).tag(sync=True)
     row_count = traitlets.Int(allow_none=True, default_value=None).tag(sync=True)
     table_html = traitlets.Unicode("").tag(sync=True)
-    sort_column = traitlets.Unicode("").tag(sync=True)
-    sort_ascending = traitlets.Bool(True).tag(sync=True)
+    sort_columns = traitlets.List(traitlets.Unicode(), []).tag(sync=True)
+    sort_ascending = traitlets.List(traitlets.Bool(), []).tag(sync=True)
     orderable_columns = traitlets.List(traitlets.Unicode(), []).tag(sync=True)
     _initial_load_complete = traitlets.Bool(False).tag(sync=True)
     _batches: Optional[blocks.PandasBatches] = None
@@ -280,23 +280,22 @@ def _set_table_html(self) -> None:
 
             # Apply sorting if a column is selected
             df_to_display = self._dataframe
-            if self.sort_column:
+            if self.sort_columns:
                 # TODO(b/463715504): Support sorting by index columns.
                 df_to_display = df_to_display.sort_values(
-                    by=self.sort_column, ascending=self.sort_ascending
+                    by=list(self.sort_columns), ascending=list(self.sort_ascending)
                 )
 
             # Reset batches when sorting changes
-            if self._last_sort_state != _SortState(
-                self.sort_column, self.sort_ascending
-            ):
+            current_sort_state = _SortState(
+                tuple(self.sort_columns), tuple(self.sort_ascending)
+            )
+            if self._last_sort_state != current_sort_state:
                 self._batches = df_to_display.to_pandas_batches(
                     page_size=self.page_size
                 )
                 self._reset_batch_cache()
-                self._last_sort_state = _SortState(
-                    self.sort_column, self.sort_ascending
-                )
+                self._last_sort_state = current_sort_state
                 if self.page != 0:
                     new_page = 0  # Reset to first page
 
@@ -348,14 +347,15 @@ def _set_table_html(self) -> None:
                 self.table_html = bigframes.display.html.render_html(
                     dataframe=page_data,
                     table_id=f"table-{self._table_id}",
+                    orderable_columns=self.orderable_columns,
                 )
 
         if new_page is not None:
             # Navigate to the new page. This triggers the observer, which will
             # re-enter _set_table_html. Since we've released the lock, this is safe.
             self.page = new_page
 
-    @traitlets.observe("sort_column", "sort_ascending")
+    @traitlets.observe("sort_columns", "sort_ascending")
     def _sort_changed(self, _change: dict[str, Any]):
         """Handler for when sorting parameters change from the frontend."""
         self._set_table_html()
@@ -375,8 +375,8 @@ def _page_size_changed(self, _change: dict[str, Any]) -> None:
         # Reset the page to 0 when page size changes to avoid invalid page states
         self.page = 0
         # Reset the sort state to default (no sort)
-        self.sort_column = ""
-        self.sort_ascending = True
+        self.sort_columns = []
+        self.sort_ascending = []
 
         # Reset batches to use new page size for future data fetching
         self._reset_batches_for_new_page_size()

bigframes/display/table_widget.js

@@ -21,7 +21,7 @@ const ModelProperty = {
 	PAGE_SIZE: "page_size",
 	ROW_COUNT: "row_count",
 	SORT_ASCENDING: "sort_ascending",
-	SORT_COLUMN: "sort_column",
+	SORT_COLUMNS: "sort_columns",
 	TABLE_HTML: "table_html",
 };
 
@@ -141,7 +141,7 @@ function render({ model, el }) {
 
 		// Get sortable columns from backend
 		const sortableColumns = model.get(ModelProperty.ORDERABLE_COLUMNS);
-		const currentSortColumn = model.get(ModelProperty.SORT_COLUMN);
+		const currentSortColumns = model.get(ModelProperty.SORT_COLUMNS);
 		const currentSortAscending = model.get(ModelProperty.SORT_ASCENDING);
 
 		// Add click handlers to column headers for sorting
@@ -161,8 +161,10 @@ function render({ model, el }) {
 
 				// Determine sort indicator and initial visibility
 				let indicator = "●"; // Default: unsorted (dot)
-				if (currentSortColumn === columnName) {
-					indicator = currentSortAscending ? "▲" : "▼";
+				const sortIndex = currentSortColumns.indexOf(columnName);
+				if (sortIndex !== -1) {
+					const isAscending = currentSortAscending[sortIndex];
+					indicator = isAscending ? "▲" : "▼";
 					indicatorSpan.style.visibility = "visible"; // Sorted arrows always visible
 				} else {
 					indicatorSpan.style.visibility = "hidden"; // Unsorted dot hidden by default
@@ -178,32 +180,57 @@ function render({ model, el }) {
 
 				// Add hover effects for unsorted columns only
 				header.addEventListener("mouseover", () => {
-					if (currentSortColumn !== columnName) {
+					if (currentSortColumns.indexOf(columnName) === -1) {
 						indicatorSpan.style.visibility = "visible";
 					}
 				});
 				header.addEventListener("mouseout", () => {
-					if (currentSortColumn !== columnName) {
+					if (currentSortColumns.indexOf(columnName) === -1) {
 						indicatorSpan.style.visibility = "hidden";
 					}
 				});
 
 				// Add click handler for three-state toggle
-				header.addEventListener(Event.CLICK, () => {
-					if (currentSortColumn === columnName) {
-						if (currentSortAscending) {
-							// Currently ascending → switch to descending
-							model.set(ModelProperty.SORT_ASCENDING, false);
+				header.addEventListener(Event.CLICK, (event) => {
+					const sortIndex = currentSortColumns.indexOf(columnName);
+					let newColumns = [...currentSortColumns];
+					let newAscending = [...currentSortAscending];
+
+					if (event.shiftKey) {
+						if (sortIndex !== -1) {
+							// Already sorted. Toggle or Remove.
+							if (newAscending[sortIndex]) {
+								// Asc -> Desc
+								newAscending[sortIndex] = false;
+							} else {
+								// Desc -> Remove
+								newColumns.splice(sortIndex, 1);
+								newAscending.splice(sortIndex, 1);
+							}
 						} else {
-							// Currently descending → clear sort (back to unsorted)
-							model.set(ModelProperty.SORT_COLUMN, "");
-							model.set(ModelProperty.SORT_ASCENDING, true);
+							// Not sorted -> Append Asc
+							newColumns.push(columnName);
+							newAscending.push(true);
 						}
 					} else {
-						// Not currently sorted → sort ascending
-						model.set(ModelProperty.SORT_COLUMN, columnName);
-						model.set(ModelProperty.SORT_ASCENDING, true);
+						// No shift key. Single column mode.
+						if (sortIndex !== -1 && newColumns.length === 1) {
+							// Already only this column. Toggle or Remove.
+							if (newAscending[sortIndex]) {
+								newAscending[sortIndex] = false;
+							} else {
+								newColumns = [];
+								newAscending = [];
+							}
+						} else {
+							// Start fresh with this column
+							newColumns = [columnName];
+							newAscending = [true];
+						}
 					}
+
+					model.set(ModelProperty.SORT_COLUMNS, newColumns);
+					model.set(ModelProperty.SORT_ASCENDING, newAscending);
 					model.save_changes();
 				});
 			}

tests/unit/display/test_anywidget.py

@@ -78,3 +78,80 @@ def handler(signum, frame):
     finally:
         if has_sigalrm:
             signal.alarm(0)
+
+
+@pytest.fixture
+def mock_df():
+    df = mock.create_autospec(bigframes.dataframe.DataFrame, instance=True)
+    df.columns = ["col1", "col2"]
+    df.dtypes = {"col1": "int64", "col2": "int64"}
+
+    mock_block = mock.Mock()
+    mock_block.has_index = False
+    df._block = mock_block
+
+    # Mock to_pandas_batches to return empty iterator or simple data
+    batch_df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+    batches = mock.MagicMock()
+    batches.__iter__.return_value = iter([batch_df])
+    batches.total_rows = 2
+    df.to_pandas_batches.return_value = batches
+
+    # Mock sort_values to return self (for chaining)
+    df.sort_values.return_value = df
+
+    return df
+
+
+def test_sorting_single_column(mock_df):
+    from bigframes.display.anywidget import TableWidget
+
+    with bigframes.option_context("display.repr_mode", "anywidget"):
+        widget = TableWidget(mock_df)
+
+    # Verify initial state
+    assert widget.sort_columns == []
+    assert widget.sort_ascending == []
+
+    # Apply sort
+    widget.sort_columns = ["col1"]
+    widget.sort_ascending = [True]
+
+    # This should trigger _sort_changed -> _set_table_html
+    # which calls df.sort_values
+
+    mock_df.sort_values.assert_called_with(by=["col1"], ascending=[True])
+
+
+def test_sorting_multi_column(mock_df):
+    from bigframes.display.anywidget import TableWidget
+
+    with bigframes.option_context("display.repr_mode", "anywidget"):
+        widget = TableWidget(mock_df)
+
+    # Apply multi-column sort
+    widget.sort_columns = ["col1", "col2"]
+    widget.sort_ascending = [True, False]
+
+    mock_df.sort_values.assert_called_with(by=["col1", "col2"], ascending=[True, False])
+
+
+def test_page_size_change_resets_sort(mock_df):
+    from bigframes.display.anywidget import TableWidget
+
+    with bigframes.option_context("display.repr_mode", "anywidget"):
+        widget = TableWidget(mock_df)
+
+    # Set sort state
+    widget.sort_columns = ["col1"]
+    widget.sort_ascending = [True]
+
+    # Change page size
+    widget.page_size = 50
+
+    # Sort should be reset
+    assert widget.sort_columns == []
+    assert widget.sort_ascending == []
+
+    # to_pandas_batches called again (reset)
+    assert mock_df.to_pandas_batches.call_count >= 2