fix: use snake_case for sorting traitlets without sync_key workaround

Author: shuoweil

bigframes/display/anywidget.py

@@ -70,7 +70,8 @@ class TableWidget(_WIDGET_BASE):
     max_columns = traitlets.Int(allow_none=True, default_value=None).tag(sync=True)
     row_count = traitlets.Int(allow_none=True, default_value=None).tag(sync=True)
     table_html = traitlets.Unicode("").tag(sync=True)
-    sort_context = traitlets.List(traitlets.Dict(), default_value=[]).tag(sync=True)
+    sort_columns = traitlets.List(traitlets.Unicode(), default_value=[]).tag(sync=True)
+    sort_ascending = traitlets.List(traitlets.Bool(), default_value=[]).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
@@ -305,8 +306,8 @@ def _set_table_html(self) -> None:
 
             # Apply sorting if a column is selected
             df_to_display = self._dataframe
-            sort_columns = [item["column"] for item in self.sort_context]
-            sort_ascending = [item["ascending"] for item in self.sort_context]
+            sort_columns = list(self.sort_columns)
+            sort_ascending = list(self.sort_ascending)
 
             if sort_columns:
                 # TODO(b/463715504): Support sorting by index columns.
@@ -382,10 +383,11 @@ def _set_table_html(self) -> None:
             # re-enter _set_table_html. Since we've released the lock, this is safe.
             self.page = new_page
 
-    @traitlets.observe("sort_context")
-    def _sort_changed(self, _change: dict[str, Any]):
+    @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()
+        if len(self.sort_columns) == len(self.sort_ascending):
+            self._set_table_html()
 
     @traitlets.observe("page")
     def _page_changed(self, _change: dict[str, Any]) -> None:
@@ -402,7 +404,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_context = []
+        self.sort_ascending = []
+        self.sort_columns = []
 
         # Reset batches to use new page size for future data fetching
         self._reset_batches_for_new_page_size()

bigframes/display/table_widget.js

@@ -20,7 +20,8 @@ const ModelProperty = {
   PAGE: 'page',
   PAGE_SIZE: 'page_size',
   ROW_COUNT: 'row_count',
-  SORT_CONTEXT: 'sort_context',
+  SORT_COLUMNS: 'sort_columns',
+  SORT_ASCENDING: 'sort_ascending',
   TABLE_HTML: 'table_html',
   MAX_COLUMNS: 'max_columns',
 };
@@ -192,10 +193,10 @@ function render({ model, el }) {
     }, 0);
 
     const sortableColumns = model.get(ModelProperty.ORDERABLE_COLUMNS);
-    const currentSortContext = model.get(ModelProperty.SORT_CONTEXT) || [];
+    const currentSortColumns = model.get(ModelProperty.SORT_COLUMNS) || [];
+    const currentSortAscending = model.get(ModelProperty.SORT_ASCENDING) || [];
 
-    const getSortIndex = (colName) =>
-      currentSortContext.findIndex((item) => item.column === colName);
+    const getSortIndex = (colName) => currentSortColumns.indexOf(colName);
 
     const headers = tableContainer.querySelectorAll('th');
     headers.forEach((header) => {
@@ -214,7 +215,7 @@ function render({ model, el }) {
         const sortIndex = getSortIndex(columnName);
 
         if (sortIndex !== -1) {
-          const isAscending = currentSortContext[sortIndex].ascending;
+          const isAscending = currentSortAscending[sortIndex];
           indicator = isAscending ? '▲' : '▼';
           indicatorSpan.style.visibility = 'visible'; // Sorted arrows always visible
         } else {
@@ -239,48 +240,46 @@ function render({ model, el }) {
           }
         });
 
-        // Add click handler for three-state toggle
         header.addEventListener(Event.CLICK, (event) => {
           const sortIndex = getSortIndex(columnName);
-          let newContext = [...currentSortContext];
+          let newSortColumns = [...currentSortColumns];
+          let newSortAscending = [...currentSortAscending];
 
           if (event.shiftKey) {
             if (sortIndex !== -1) {
               // Already sorted. Toggle or Remove.
-              if (newContext[sortIndex].ascending) {
+              if (newSortAscending[sortIndex]) {
                 // Asc -> Desc
-                // Clone object to avoid mutation issues
-                newContext[sortIndex] = {
-                  ...newContext[sortIndex],
-                  ascending: false,
-                };
+                newSortAscending[sortIndex] = false;
               } else {
                 // Desc -> Remove
-                newContext.splice(sortIndex, 1);
+                newSortColumns.splice(sortIndex, 1);
+                newSortAscending.splice(sortIndex, 1);
               }
             } else {
               // Not sorted -> Append Asc
-              newContext.push({ column: columnName, ascending: true });
+              newSortColumns.push(columnName);
+              newSortAscending.push(true);
             }
           } else {
             // No shift key. Single column mode.
-            if (sortIndex !== -1 && newContext.length === 1) {
+            if (sortIndex !== -1 && newSortColumns.length === 1) {
               // Already only this column. Toggle or Remove.
-              if (newContext[sortIndex].ascending) {
-                newContext[sortIndex] = {
-                  ...newContext[sortIndex],
-                  ascending: false,
-                };
+              if (newSortAscending[sortIndex]) {
+                newSortAscending[sortIndex] = false;
               } else {
-                newContext = [];
+                newSortColumns = [];
+                newSortAscending = [];
               }
             } else {
               // Start fresh with this column
-              newContext = [{ column: columnName, ascending: true }];
+              newSortColumns = [columnName];
+              newSortAscending = [true];
             }
           }
 
-          model.set(ModelProperty.SORT_CONTEXT, newContext);
+          model.set(ModelProperty.SORT_ASCENDING, newSortAscending);
+          model.set(ModelProperty.SORT_COLUMNS, newSortColumns);
           model.save_changes();
         });
       }

tests/unit/display/test_anywidget.py

@@ -134,10 +134,12 @@ def test_sorting_single_column(mock_df):
         widget = TableWidget(mock_df)
 
     # Verify initial state
-    assert widget.sort_context == []
+    assert widget.sort_columns == []
+    assert widget.sort_ascending == []
 
     # Apply sort
-    widget.sort_context = [{"column": "col1", "ascending": True}]
+    widget.sort_columns = ["col1"]
+    widget.sort_ascending = [True]
 
     # This should trigger _sort_changed -> _set_table_html
     # which calls df.sort_values
@@ -153,10 +155,8 @@ def test_sorting_multi_column(mock_df):
         widget = TableWidget(mock_df)
 
     # Apply multi-column sort
-    widget.sort_context = [
-        {"column": "col1", "ascending": True},
-        {"column": "col2", "ascending": False},
-    ]
+    widget.sort_columns = ["col1", "col2"]
+    widget.sort_ascending = [True, False]
 
     mock_df.sort_values.assert_called_with(by=["col1", "col2"], ascending=[True, False])
 
@@ -169,13 +169,15 @@ def test_page_size_change_resets_sort(mock_df):
         widget = TableWidget(mock_df)
 
     # Set sort state
-    widget.sort_context = [{"column": "col1", "ascending": True}]
+    widget.sort_columns = ["col1"]
+    widget.sort_ascending = [True]
 
     # Change page size
     widget.page_size = 50
 
     # Sort should be reset
-    assert widget.sort_context == []
+    assert widget.sort_columns == []
+    assert widget.sort_ascending == []
 
     # to_pandas_batches called again (reset)
     assert mock_df.to_pandas_batches.call_count >= 2