Skip to content

Commit b755c55

Browse files
shuoweilshuoweil
committed
refactor(anywidget): merge sort_columns/ascending into sort_context
1 parent b5306c1 commit b755c55

File tree

5 files changed

+73
-90
lines changed

5 files changed

+73
-90
lines changed

‎bigframes/display/anywidget.py‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,7 @@ class TableWidget(_WIDGET_BASE):
6868
page_size = traitlets.Int(0).tag(sync=True)
6969
row_count = traitlets.Int(allow_none=True, default_value=None).tag(sync=True)
7070
table_html = traitlets.Unicode("").tag(sync=True)
71-
sort_columns = traitlets.List(traitlets.Unicode(), []).tag(sync=True)
72-
sort_ascending = traitlets.List(traitlets.Bool(), []).tag(sync=True)
71+
sort_context = traitlets.List(traitlets.Dict(), default_value=[]).tag(sync=True)
7372
orderable_columns = traitlets.List(traitlets.Unicode(), []).tag(sync=True)
7473
_initial_load_complete = traitlets.Bool(False).tag(sync=True)
7574
_batches: Optional[blocks.PandasBatches] = None
@@ -280,16 +279,17 @@ def _set_table_html(self) -> None:
280279

281280
# Apply sorting if a column is selected
282281
df_to_display = self._dataframe
283-
if self.sort_columns:
282+
sort_columns = [item["column"] for item in self.sort_context]
283+
sort_ascending = [item["ascending"] for item in self.sort_context]
284+
285+
if sort_columns:
284286
# TODO(b/463715504): Support sorting by index columns.
285287
df_to_display = df_to_display.sort_values(
286-
by=list(self.sort_columns), ascending=list(self.sort_ascending)
288+
by=sort_columns, ascending=sort_ascending
287289
)
288290

289291
# Reset batches when sorting changes
290-
current_sort_state = _SortState(
291-
tuple(self.sort_columns), tuple(self.sort_ascending)
292-
)
292+
current_sort_state = _SortState(tuple(sort_columns), tuple(sort_ascending))
293293
if self._last_sort_state != current_sort_state:
294294
self._batches = df_to_display.to_pandas_batches(
295295
page_size=self.page_size
@@ -355,7 +355,7 @@ def _set_table_html(self) -> None:
355355
# re-enter _set_table_html. Since we've released the lock, this is safe.
356356
self.page = new_page
357357

358-
@traitlets.observe("sort_columns", "sort_ascending")
358+
@traitlets.observe("sort_context")
359359
def _sort_changed(self, _change: dict[str, Any]):
360360
"""Handler for when sorting parameters change from the frontend."""
361361
self._set_table_html()
@@ -375,8 +375,7 @@ def _page_size_changed(self, _change: dict[str, Any]) -> None:
375375
# Reset the page to 0 when page size changes to avoid invalid page states
376376
self.page = 0
377377
# Reset the sort state to default (no sort)
378-
self.sort_columns = []
379-
self.sort_ascending = []
378+
self.sort_context = []
380379

381380
# Reset batches to use new page size for future data fetching
382381
self._reset_batches_for_new_page_size()

‎bigframes/display/table_widget.js‎

Lines changed: 29 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,7 @@ const ModelProperty = {
2020
PAGE: "page",
2121
PAGE_SIZE: "page_size",
2222
ROW_COUNT: "row_count",
23-
SORT_ASCENDING: "sort_ascending",
24-
SORT_COLUMNS: "sort_columns",
23+
SORT_CONTEXT: "sort_context",
2524
TABLE_HTML: "table_html",
2625
};
2726

@@ -141,8 +140,10 @@ function render({ model, el }) {
141140

142141
// Get sortable columns from backend
143142
const sortableColumns = model.get(ModelProperty.ORDERABLE_COLUMNS);
144-
const currentSortColumns = model.get(ModelProperty.SORT_COLUMNS);
145-
const currentSortAscending = model.get(ModelProperty.SORT_ASCENDING);
143+
const currentSortContext = model.get(ModelProperty.SORT_CONTEXT) || [];
144+
145+
const getSortIndex = (colName) =>
146+
currentSortContext.findIndex((item) => item.column === colName);
146147

147148
// Add click handlers to column headers for sorting
148149
const headers = tableContainer.querySelectorAll("th");
@@ -161,9 +162,10 @@ function render({ model, el }) {
161162

162163
// Determine sort indicator and initial visibility
163164
let indicator = "●"; // Default: unsorted (dot)
164-
const sortIndex = currentSortColumns.indexOf(columnName);
165+
const sortIndex = getSortIndex(columnName);
166+
165167
if (sortIndex !== -1) {
166-
const isAscending = currentSortAscending[sortIndex];
168+
const isAscending = currentSortContext[sortIndex].ascending;
167169
indicator = isAscending ? "▲" : "▼";
168170
indicatorSpan.style.visibility = "visible"; // Sorted arrows always visible
169171
} else {
@@ -180,57 +182,58 @@ function render({ model, el }) {
180182

181183
// Add hover effects for unsorted columns only
182184
header.addEventListener("mouseover", () => {
183-
if (currentSortColumns.indexOf(columnName) === -1) {
185+
if (getSortIndex(columnName) === -1) {
184186
indicatorSpan.style.visibility = "visible";
185187
}
186188
});
187189
header.addEventListener("mouseout", () => {
188-
if (currentSortColumns.indexOf(columnName) === -1) {
190+
if (getSortIndex(columnName) === -1) {
189191
indicatorSpan.style.visibility = "hidden";
190192
}
191193
});
192194

193195
// Add click handler for three-state toggle
194196
header.addEventListener(Event.CLICK, (event) => {
195-
const sortIndex = currentSortColumns.indexOf(columnName);
196-
let newColumns = [...currentSortColumns];
197-
let newAscending = [...currentSortAscending];
197+
const sortIndex = getSortIndex(columnName);
198+
let newContext = [...currentSortContext];
198199

199200
if (event.shiftKey) {
200201
if (sortIndex !== -1) {
201202
// Already sorted. Toggle or Remove.
202-
if (newAscending[sortIndex]) {
203+
if (newContext[sortIndex].ascending) {
203204
// Asc -> Desc
204-
newAscending[sortIndex] = false;
205+
// Clone object to avoid mutation issues
206+
newContext[sortIndex] = {
207+
...newContext[sortIndex],
208+
ascending: false,
209+
};
205210
} else {
206211
// Desc -> Remove
207-
newColumns.splice(sortIndex, 1);
208-
newAscending.splice(sortIndex, 1);
212+
newContext.splice(sortIndex, 1);
209213
}
210214
} else {
211215
// Not sorted -> Append Asc
212-
newColumns.push(columnName);
213-
newAscending.push(true);
216+
newContext.push({ column: columnName, ascending: true });
214217
}
215218
} else {
216219
// No shift key. Single column mode.
217-
if (sortIndex !== -1 && newColumns.length === 1) {
220+
if (sortIndex !== -1 && newContext.length === 1) {
218221
// Already only this column. Toggle or Remove.
219-
if (newAscending[sortIndex]) {
220-
newAscending[sortIndex] = false;
222+
if (newContext[sortIndex].ascending) {
223+
newContext[sortIndex] = {
224+
...newContext[sortIndex],
225+
ascending: false,
226+
};
221227
} else {
222-
newColumns = [];
223-
newAscending = [];
228+
newContext = [];
224229
}
225230
} else {
226231
// Start fresh with this column
227-
newColumns = [columnName];
228-
newAscending = [true];
232+
newContext = [{ column: columnName, ascending: true }];
229233
}
230234
}
231235

232-
model.set(ModelProperty.SORT_COLUMNS, newColumns);
233-
model.set(ModelProperty.SORT_ASCENDING, newAscending);
236+
model.set(ModelProperty.SORT_CONTEXT, newContext);
234237
model.save_changes();
235238
});
236239
}

‎tests/js/table_widget.test.js‎

Lines changed: 20 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -83,10 +83,7 @@ describe("TableWidget", () => {
8383
if (property === "orderable_columns") {
8484
return ["col1"];
8585
}
86-
if (property === "sort_columns") {
87-
return [];
88-
}
89-
if (property === "sort_ascending") {
86+
if (property === "sort_context") {
9087
return [];
9188
}
9289
return null;
@@ -103,8 +100,9 @@ describe("TableWidget", () => {
103100
const header = el.querySelector("th");
104101
header.click();
105102

106-
expect(model.set).toHaveBeenCalledWith("sort_columns", ["col1"]);
107-
expect(model.set).toHaveBeenCalledWith("sort_ascending", [true]);
103+
expect(model.set).toHaveBeenCalledWith("sort_context", [
104+
{ column: "col1", ascending: true },
105+
]);
108106
expect(model.save_changes).toHaveBeenCalled();
109107
});
110108

@@ -117,11 +115,8 @@ describe("TableWidget", () => {
117115
if (property === "orderable_columns") {
118116
return ["col1"];
119117
}
120-
if (property === "sort_columns") {
121-
return ["col1"];
122-
}
123-
if (property === "sort_ascending") {
124-
return [true];
118+
if (property === "sort_context") {
119+
return [{ column: "col1", ascending: true }];
125120
}
126121
return null;
127122
});
@@ -137,7 +132,9 @@ describe("TableWidget", () => {
137132
const header = el.querySelector("th");
138133
header.click();
139134

140-
expect(model.set).toHaveBeenCalledWith("sort_ascending", [false]);
135+
expect(model.set).toHaveBeenCalledWith("sort_context", [
136+
{ column: "col1", ascending: false },
137+
]);
141138
expect(model.save_changes).toHaveBeenCalled();
142139
});
143140

@@ -150,11 +147,8 @@ describe("TableWidget", () => {
150147
if (property === "orderable_columns") {
151148
return ["col1"];
152149
}
153-
if (property === "sort_columns") {
154-
return ["col1"];
155-
}
156-
if (property === "sort_ascending") {
157-
return [false];
150+
if (property === "sort_context") {
151+
return [{ column: "col1", ascending: false }];
158152
}
159153
return null;
160154
});
@@ -170,8 +164,7 @@ describe("TableWidget", () => {
170164
const header = el.querySelector("th");
171165
header.click();
172166

173-
expect(model.set).toHaveBeenCalledWith("sort_columns", []);
174-
expect(model.set).toHaveBeenCalledWith("sort_ascending", []);
167+
expect(model.set).toHaveBeenCalledWith("sort_context", []);
175168
expect(model.save_changes).toHaveBeenCalled();
176169
});
177170

@@ -184,11 +177,8 @@ describe("TableWidget", () => {
184177
if (property === "orderable_columns") {
185178
return ["col1", "col2"];
186179
}
187-
if (property === "sort_columns") {
188-
return ["col1"];
189-
}
190-
if (property === "sort_ascending") {
191-
return [true];
180+
if (property === "sort_context") {
181+
return [{ column: "col1", ascending: true }];
192182
}
193183
return null;
194184
});
@@ -218,11 +208,8 @@ describe("TableWidget", () => {
218208
if (property === "orderable_columns") {
219209
return ["col1", "col2"];
220210
}
221-
if (property === "sort_columns") {
222-
return ["col1"];
223-
}
224-
if (property === "sort_ascending") {
225-
return [true];
211+
if (property === "sort_context") {
212+
return [{ column: "col1", ascending: true }];
226213
}
227214
return null;
228215
});
@@ -246,8 +233,10 @@ describe("TableWidget", () => {
246233
});
247234
header2.dispatchEvent(clickEvent);
248235

249-
expect(model.set).toHaveBeenCalledWith("sort_columns", ["col1", "col2"]);
250-
expect(model.set).toHaveBeenCalledWith("sort_ascending", [true, true]);
236+
expect(model.set).toHaveBeenCalledWith("sort_context", [
237+
{ column: "col1", ascending: true },
238+
{ column: "col2", ascending: true },
239+
]);
251240
expect(model.save_changes).toHaveBeenCalled();
252241
});
253242
});

‎tests/system/small/test_anywidget.py‎

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -774,8 +774,7 @@ def test_widget_sort_should_sort_ascending_on_first_click(
774774
Given a widget, when a column header is clicked for the first time,
775775
then the data should be sorted by that column in ascending order.
776776
"""
777-
table_widget.sort_ascending = [True]
778-
table_widget.sort_columns = ["id"]
777+
table_widget.sort_context = [{"column": "id", "ascending": True}]
779778

780779
expected_slice = paginated_pandas_df.sort_values("id", ascending=True).iloc[0:2]
781780
html = table_widget.table_html
@@ -790,11 +789,10 @@ def test_widget_sort_should_sort_descending_on_second_click(
790789
Given a widget sorted by a column, when the same column header is clicked again,
791790
then the data should be sorted by that column in descending order.
792791
"""
793-
table_widget.sort_ascending = [True]
794-
table_widget.sort_columns = ["id"]
792+
table_widget.sort_context = [{"column": "id", "ascending": True}]
795793

796794
# Second click
797-
table_widget.sort_ascending = [False]
795+
table_widget.sort_context = [{"column": "id", "ascending": False}]
798796

799797
expected_slice = paginated_pandas_df.sort_values("id", ascending=False).iloc[0:2]
800798
html = table_widget.table_html
@@ -809,12 +807,10 @@ def test_widget_sort_should_switch_column_and_sort_ascending(
809807
Given a widget sorted by a column, when a different column header is clicked,
810808
then the data should be sorted by the new column in ascending order.
811809
"""
812-
table_widget.sort_ascending = [True]
813-
table_widget.sort_columns = ["id"]
810+
table_widget.sort_context = [{"column": "id", "ascending": True}]
814811

815812
# Click on a different column
816-
table_widget.sort_ascending = [True]
817-
table_widget.sort_columns = ["value"]
813+
table_widget.sort_context = [{"column": "value", "ascending": True}]
818814

819815
expected_slice = paginated_pandas_df.sort_values("value", ascending=True).iloc[0:2]
820816
html = table_widget.table_html
@@ -829,8 +825,7 @@ def test_widget_sort_should_be_maintained_after_pagination(
829825
Given a sorted widget, when the user navigates to the next page,
830826
then the sorting should be maintained.
831827
"""
832-
table_widget.sort_ascending = [True]
833-
table_widget.sort_columns = ["id"]
828+
table_widget.sort_context = [{"column": "id", "ascending": True}]
834829

835830
# Go to the second page
836831
table_widget.page = 1
@@ -848,8 +843,7 @@ def test_widget_sort_should_reset_on_page_size_change(
848843
Given a sorted widget, when the page size is changed,
849844
then the sorting should be reset.
850845
"""
851-
table_widget.sort_ascending = [True]
852-
table_widget.sort_columns = ["id"]
846+
table_widget.sort_context = [{"column": "id", "ascending": True}]
853847

854848
table_widget.page_size = 3
855849

‎tests/unit/display/test_anywidget.py‎

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -110,12 +110,10 @@ def test_sorting_single_column(mock_df):
110110
widget = TableWidget(mock_df)
111111

112112
# Verify initial state
113-
assert widget.sort_columns == []
114-
assert widget.sort_ascending == []
113+
assert widget.sort_context == []
115114

116115
# Apply sort
117-
widget.sort_columns = ["col1"]
118-
widget.sort_ascending = [True]
116+
widget.sort_context = [{"column": "col1", "ascending": True}]
119117

120118
# This should trigger _sort_changed -> _set_table_html
121119
# which calls df.sort_values
@@ -130,8 +128,10 @@ def test_sorting_multi_column(mock_df):
130128
widget = TableWidget(mock_df)
131129

132130
# Apply multi-column sort
133-
widget.sort_columns = ["col1", "col2"]
134-
widget.sort_ascending = [True, False]
131+
widget.sort_context = [
132+
{"column": "col1", "ascending": True},
133+
{"column": "col2", "ascending": False},
134+
]
135135

136136
mock_df.sort_values.assert_called_with(by=["col1", "col2"], ascending=[True, False])
137137

@@ -143,15 +143,13 @@ def test_page_size_change_resets_sort(mock_df):
143143
widget = TableWidget(mock_df)
144144

145145
# Set sort state
146-
widget.sort_columns = ["col1"]
147-
widget.sort_ascending = [True]
146+
widget.sort_context = [{"column": "col1", "ascending": True}]
148147

149148
# Change page size
150149
widget.page_size = 50
151150

152151
# Sort should be reset
153-
assert widget.sort_columns == []
154-
assert widget.sort_ascending == []
152+
assert widget.sort_context == []
155153

156154
# to_pandas_batches called again (reset)
157155
assert mock_df.to_pandas_batches.call_count >= 2

0 commit comments

Comments
 (0)