diff --git a/bigframes/display/anywidget.py b/bigframes/display/anywidget.py index 40d04a1d71..ea656310b5 100644 --- a/bigframes/display/anywidget.py +++ b/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() diff --git a/bigframes/display/table_widget.js b/bigframes/display/table_widget.js index 314bf771d0..582c98bfd9 100644 --- a/bigframes/display/table_widget.js +++ b/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(); }); } diff --git a/tests/js/table_widget.test.js b/tests/js/table_widget.test.js index d701d8692e..283fb5c726 100644 --- a/tests/js/table_widget.test.js +++ b/tests/js/table_widget.test.js @@ -82,7 +82,10 @@ describe('TableWidget', () => { if (property === 'orderable_columns') { return ['col1']; } - if (property === 'sort_context') { + if (property === 'sort_columns') { + return []; + } + if (property === 'sort_ascending') { return []; } return null; @@ -99,9 +102,8 @@ describe('TableWidget', () => { const header = el.querySelector('th'); header.click(); - expect(model.set).toHaveBeenCalledWith('sort_context', [ - { column: 'col1', ascending: true }, - ]); + expect(model.set).toHaveBeenCalledWith('sort_ascending', [true]); + expect(model.set).toHaveBeenCalledWith('sort_columns', ['col1']); expect(model.save_changes).toHaveBeenCalled(); }); @@ -114,8 +116,11 @@ describe('TableWidget', () => { if (property === 'orderable_columns') { return ['col1']; } - if (property === 'sort_context') { - return [{ column: 'col1', ascending: true }]; + if (property === 'sort_columns') { + return ['col1']; + } + if (property === 'sort_ascending') { + return [true]; } return null; }); @@ -131,9 +136,8 @@ describe('TableWidget', () => { const header = el.querySelector('th'); header.click(); - expect(model.set).toHaveBeenCalledWith('sort_context', [ - { column: 'col1', ascending: false }, - ]); + expect(model.set).toHaveBeenCalledWith('sort_ascending', [false]); + expect(model.set).toHaveBeenCalledWith('sort_columns', ['col1']); expect(model.save_changes).toHaveBeenCalled(); }); @@ -146,8 +150,11 @@ describe('TableWidget', () => { if (property === 'orderable_columns') { return ['col1']; } - if (property === 'sort_context') { - return [{ column: 'col1', ascending: false }]; + if (property === 'sort_columns') { + return ['col1']; + } + if (property === 'sort_ascending') { + return [false]; } return null; }); @@ -163,7 +170,8 @@ describe('TableWidget', () => { const header = el.querySelector('th'); header.click(); - expect(model.set).toHaveBeenCalledWith('sort_context', []); + expect(model.set).toHaveBeenCalledWith('sort_ascending', []); + expect(model.set).toHaveBeenCalledWith('sort_columns', []); expect(model.save_changes).toHaveBeenCalled(); }); @@ -176,8 +184,11 @@ describe('TableWidget', () => { if (property === 'orderable_columns') { return ['col1', 'col2']; } - if (property === 'sort_context') { - return [{ column: 'col1', ascending: true }]; + if (property === 'sort_columns') { + return ['col1']; + } + if (property === 'sort_ascending') { + return [true]; } return null; }); @@ -207,8 +218,11 @@ describe('TableWidget', () => { if (property === 'orderable_columns') { return ['col1', 'col2']; } - if (property === 'sort_context') { - return [{ column: 'col1', ascending: true }]; + if (property === 'sort_columns') { + return ['col1']; + } + if (property === 'sort_ascending') { + return [true]; } return null; }); @@ -232,10 +246,8 @@ describe('TableWidget', () => { }); header2.dispatchEvent(clickEvent); - expect(model.set).toHaveBeenCalledWith('sort_context', [ - { column: 'col1', ascending: true }, - { column: 'col2', ascending: true }, - ]); + expect(model.set).toHaveBeenCalledWith('sort_ascending', [true, true]); + expect(model.set).toHaveBeenCalledWith('sort_columns', ['col1', 'col2']); expect(model.save_changes).toHaveBeenCalled(); }); }); @@ -418,7 +430,10 @@ describe('TableWidget', () => { // Only actual columns are orderable return ['col1', 'col10']; } - if (property === 'sort_context') { + if (property === 'sort_columns') { + return []; + } + if (property === 'sort_ascending') { return []; } return null; diff --git a/tests/unit/display/test_anywidget.py b/tests/unit/display/test_anywidget.py index d8c8c64ceb..dd8aea4a90 100644 --- a/tests/unit/display/test_anywidget.py +++ b/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