Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions bigframes/display/anywidget.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand All @@ -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()
Expand Down
47 changes: 23 additions & 24 deletions bigframes/display/table_widget.js
Original file line number Diff line number Diff line change
Expand Up @@ -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',
};
Expand Down Expand Up @@ -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) => {
Expand All @@ -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 {
Expand All @@ -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();
});
}
Expand Down
57 changes: 36 additions & 21 deletions tests/js/table_widget.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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();
});

Expand All @@ -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;
});
Expand All @@ -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();
});

Expand All @@ -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;
});
Expand All @@ -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();
});

Expand All @@ -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;
});
Expand Down Expand Up @@ -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;
});
Expand All @@ -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();
});
});
Expand Down Expand Up @@ -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;
Expand Down
18 changes: 10 additions & 8 deletions tests/unit/display/test_anywidget.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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])

Expand All @@ -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
Loading