Merge branch 'main' into sycai_type_usage

Author: sycai

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,6 +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=self.max_columns,
                 )
 
         if new_page is not None:
@@ -382,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/html.py

@@ -46,21 +46,51 @@ def render_html(
     dataframe: pd.DataFrame,
     table_id: str,
     orderable_columns: list[str] | None = None,
+    max_columns: int | None = None,
 ) -> str:
     """Render a pandas DataFrame to HTML with specific styling."""
     orderable_columns = orderable_columns or []
     classes = "dataframe table table-striped table-hover"
     table_html_parts = [f'<table border="1" class="{classes}" id="{table_id}">']
-    table_html_parts.append(_render_table_header(dataframe, orderable_columns))
-    table_html_parts.append(_render_table_body(dataframe))
+
+    # Handle column truncation
+    columns = list(dataframe.columns)
+    if max_columns is not None and max_columns > 0 and len(columns) > max_columns:
+        half = max_columns // 2
+        left_columns = columns[:half]
+        # Ensure we don't take more than available if half is 0 or calculation is weird,
+        # but typical case is safe.
+        right_count = max_columns - half
+        right_columns = columns[-right_count:] if right_count > 0 else []
+        show_ellipsis = True
+    else:
+        left_columns = columns
+        right_columns = []
+        show_ellipsis = False
+
+    table_html_parts.append(
+        _render_table_header(
+            dataframe, orderable_columns, left_columns, right_columns, show_ellipsis
+        )
+    )
+    table_html_parts.append(
+        _render_table_body(dataframe, left_columns, right_columns, show_ellipsis)
+    )
     table_html_parts.append("</table>")
     return "".join(table_html_parts)
 
 
-def _render_table_header(dataframe: pd.DataFrame, orderable_columns: list[str]) -> str:
+def _render_table_header(
+    dataframe: pd.DataFrame,
+    orderable_columns: list[str],
+    left_columns: list[Any],
+    right_columns: list[Any],
+    show_ellipsis: bool,
+) -> str:
     """Render the header of the HTML table."""
     header_parts = ["  <thead>", "    <tr>"]
-    for col in dataframe.columns:
+
+    def render_col_header(col):
         th_classes = []
         if col in orderable_columns:
             th_classes.append("sortable")
@@ -69,19 +99,38 @@ def _render_table_header(dataframe: pd.DataFrame, orderable_columns: list[str])
             f'      <th {class_str}><div class="bf-header-content">'
             f"{html.escape(str(col))}</div></th>"
         )
+
+    for col in left_columns:
+        render_col_header(col)
+
+    if show_ellipsis:
+        header_parts.append(
+            '      <th><div class="bf-header-content" style="cursor: default;">...</div></th>'
+        )
+
+    for col in right_columns:
+        render_col_header(col)
+
     header_parts.extend(["    </tr>", "  </thead>"])
     return "\n".join(header_parts)
 
 
-def _render_table_body(dataframe: pd.DataFrame) -> str:
+def _render_table_body(
+    dataframe: pd.DataFrame,
+    left_columns: list[Any],
+    right_columns: list[Any],
+    show_ellipsis: bool,
+) -> str:
     """Render the body of the HTML table."""
     body_parts = ["  <tbody>"]
     precision = options.display.precision
 
     for i in range(len(dataframe)):
         body_parts.append("    <tr>")
         row = dataframe.iloc[i]
-        for col_name, value in row.items():
+
+        def render_col_cell(col_name):
+            value = row[col_name]
             dtype = dataframe.dtypes.loc[col_name]  # type: ignore
             align = "right" if _is_dtype_numeric(dtype) else "left"
 
@@ -101,6 +150,17 @@ def _render_table_body(dataframe: pd.DataFrame) -> str:
                     f'      <td class="cell-align-{align}">'
                     f"{html.escape(cell_content)}</td>"
                 )
+
+        for col in left_columns:
+            render_col_cell(col)
+
+        if show_ellipsis:
+            # Ellipsis cell
+            body_parts.append('      <td class="cell-align-left">...</td>')
+
+        for col in right_columns:
+            render_col_cell(col)
+
         body_parts.append("    </tr>")
     body_parts.append("  </tbody>")
     return "\n".join(body_parts)

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" (all columns)
+  const maxColumnOptions = [5, 10, 15, 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);

notebooks/dataframes/anywidget_mode.ipynb

@@ -51,7 +51,8 @@
     "- **Rich DataFrames & Series:** Both DataFrames and Series are displayed as interactive widgets.\n",
     "- **Pagination:** Navigate through large datasets page by page without overwhelming the output.\n",
     "- **Column Sorting:** Click column headers to toggle between ascending, descending, and unsorted views. Use **Shift + Click** to sort by multiple columns.\n",
-    "- **Column Resizing:** Drag the dividers between column headers to adjust their width."
+    "- **Column Resizing:** Drag the dividers between column headers to adjust their width.\n",
+    "- **Max Columns Control:** Limit the number of displayed columns to improve performance and readability for wide datasets."
    ]
   },
   {
@@ -511,7 +512,10 @@
    "metadata": {},
    "source": [
     "### Adjustable Column Widths\n",
-    "You can easily adjust the width of any column in the table. Simply hover your mouse over the vertical dividers between column headers. When the cursor changes to a resize icon, click and drag to expand or shrink the column to your desired width. This allows for better readability and customization of your table view."
+    "You can easily adjust the width of any column in the table. Simply hover your mouse over the vertical dividers between column headers. When the cursor changes to a resize icon, click and drag to expand or shrink the column to your desired width. This allows for better readability and customization of your table view.\n",
+    "\n",
+    "### Control Maximum Columns\n",
+    "You can control the number of columns displayed in the widget using the **Max columns** dropdown in the footer. This is useful for wide DataFrames where you want to focus on a subset of columns or improve rendering performance. Options include 3, 5, 7, 10, 20, or All."
    ]
   },
   {