feat(display): add max columns dropdown to TableWidget

Author: shuoweil

bigframes/display/anywidget.py

@@ -66,6 +66,7 @@ class TableWidget(_WIDGET_BASE):
 
     page = traitlets.Int(0).tag(sync=True)
     page_size = traitlets.Int(0).tag(sync=True)
+    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)
@@ -103,10 +104,13 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
 
         # respect display options for initial page size
         initial_page_size = bigframes.options.display.max_rows
+        initial_max_columns = bigframes.options.display.max_columns
 
         # set traitlets properties that trigger observers
         # TODO(b/462525985): Investigate and improve TableWidget UX for DataFrames with a large number of columns.
         self.page_size = initial_page_size
+        self.max_columns = initial_max_columns
+
         # TODO(b/469861913): Nested columns from structs (e.g., 'struct_col.name') are not currently sortable.
         # TODO(b/463754889): Support non-string column labels for sorting.
         if all(isinstance(col, str) for col in dataframe.columns):
@@ -218,6 +222,14 @@ def _validate_page_size(self, proposal: dict[str, Any]) -> int:
         max_page_size = 1000
         return min(value, max_page_size)
 
+    @traitlets.validate("max_columns")
+    def _validate_max_columns(self, proposal: dict[str, Any]) -> int:
+        """Validate max columns to ensure it's positive or 0 (for all)."""
+        value = proposal["value"]
+        if value is None:
+            return 0  # Normalize None to 0 for traitlet
+        return max(0, value)
+
     def _get_next_batch(self) -> bool:
         """
         Gets the next batch of data from the generator and appends to cache.
@@ -348,7 +360,7 @@ def _set_table_html(self) -> None:
                     dataframe=page_data,
                     table_id=f"table-{self._table_id}",
                     orderable_columns=self.orderable_columns,
-                    max_columns=bigframes.options.display.max_columns,
+                    max_columns=self.max_columns,
                 )
 
         if new_page is not None:
@@ -383,3 +395,10 @@ def _page_size_changed(self, _change: dict[str, Any]) -> None:
 
         # Update the table display
         self._set_table_html()
+
+    @traitlets.observe("max_columns")
+    def _max_columns_changed(self, _change: dict[str, Any]) -> None:
+        """Handler for when max columns is changed from the frontend."""
+        if not self._initial_load_complete:
+            return
+        self._set_table_html()

bigframes/display/table_widget.css

@@ -117,15 +117,24 @@ body[data-theme='dark'] .bigframes-widget.bigframes-widget {
   margin: 0 8px;
 }
 
-.bigframes-widget .page-size {
+.bigframes-widget .settings {
   align-items: center;
   display: flex;
   flex-direction: row;
-  gap: 4px;
+  gap: 16px;
   justify-content: end;
 }
 
-.bigframes-widget .page-size label {
+.bigframes-widget .page-size,
+.bigframes-widget .max-columns {
+  align-items: center;
+  display: flex;
+  flex-direction: row;
+  gap: 4px;
+}
+
+.bigframes-widget .page-size label,
+.bigframes-widget .max-columns label {
   margin-right: 8px;
 }
 

bigframes/display/table_widget.js

@@ -22,6 +22,7 @@ const ModelProperty = {
   ROW_COUNT: 'row_count',
   SORT_CONTEXT: 'sort_context',
   TABLE_HTML: 'table_html',
+  MAX_COLUMNS: 'max_columns',
 };
 
 const Event = {
@@ -71,6 +72,10 @@ function render({ model, el }) {
     attributeFilter: ['class', 'data-theme', 'data-vscode-theme-kind'],
   });
 
+  // Settings controls container
+  const settingsContainer = document.createElement('div');
+  settingsContainer.classList.add('settings');
+
   // Pagination controls
   const paginationContainer = document.createElement('div');
   paginationContainer.classList.add('pagination');
@@ -102,6 +107,32 @@ function render({ model, el }) {
     pageSizeInput.appendChild(option);
   }
 
+  // Max columns controls
+  const maxColumnsContainer = document.createElement('div');
+  maxColumnsContainer.classList.add('max-columns');
+  const maxColumnsLabel = document.createElement('label');
+  const maxColumnsInput = document.createElement('select');
+
+  maxColumnsLabel.textContent = 'Max columns:';
+
+  // 0 represents "All"
+  const maxColumnOptions = [3, 5, 7, 10, 20, 0];
+  for (const cols of maxColumnOptions) {
+    const option = document.createElement('option');
+    option.value = cols;
+    option.textContent = cols === 0 ? 'All' : cols;
+
+    const currentMax = model.get(ModelProperty.MAX_COLUMNS);
+    // Handle None/null from python as 0/All
+    const currentMaxVal =
+      currentMax === null || currentMax === undefined ? 0 : currentMax;
+
+    if (cols === currentMaxVal) {
+      option.selected = true;
+    }
+    maxColumnsInput.appendChild(option);
+  }
+
   function updateButtonStates() {
     const currentPage = model.get(ModelProperty.PAGE);
     const pageSize = model.get(ModelProperty.PAGE_SIZE);
@@ -259,6 +290,12 @@ function render({ model, el }) {
     }
   });
 
+  maxColumnsInput.addEventListener(Event.CHANGE, (e) => {
+    const newVal = Number(e.target.value);
+    model.set(ModelProperty.MAX_COLUMNS, newVal);
+    model.save_changes();
+  });
+
   model.on(Event.CHANGE_TABLE_HTML, handleTableHTMLChange);
   model.on(`change:${ModelProperty.ROW_COUNT}`, updateButtonStates);
   model.on(`change:${ModelProperty.ERROR_MESSAGE}`, handleErrorMessageChange);
@@ -270,11 +307,19 @@ function render({ model, el }) {
   paginationContainer.appendChild(prevPage);
   paginationContainer.appendChild(pageIndicator);
   paginationContainer.appendChild(nextPage);
+
   pageSizeContainer.appendChild(pageSizeLabel);
   pageSizeContainer.appendChild(pageSizeInput);
+
+  maxColumnsContainer.appendChild(maxColumnsLabel);
+  maxColumnsContainer.appendChild(maxColumnsInput);
+
+  settingsContainer.appendChild(maxColumnsContainer);
+  settingsContainer.appendChild(pageSizeContainer);
+
   footer.appendChild(rowCountLabel);
   footer.appendChild(paginationContainer);
-  footer.appendChild(pageSizeContainer);
+  footer.appendChild(settingsContainer);
 
   el.appendChild(errorContainer);
   el.appendChild(tableContainer);