feat: display custom multi index column in anywidget mode

Author: shuoweil

bigframes/dataframe.py

@@ -1969,7 +1969,7 @@ def to_pandas_batches(
         max_results: Optional[int] = None,
         *,
         allow_large_results: Optional[bool] = None,
-    ) -> Iterable[pandas.DataFrame]:
+    ) -> blocks.PandasBatches:
         """Stream DataFrame results to an iterable of pandas DataFrame.
 
         page_size and max_results determine the size and number of batches,

bigframes/display/anywidget.py

@@ -229,6 +229,10 @@ def _get_next_batch(self) -> bool:
         except StopIteration:
             self._all_data_loaded = True
             return False
+        except Exception as e:
+            # Handle other potential errors
+            self._error_message = f"Error loading data: {str(e)}"
+            return False
 
     @property
     def _batch_iterator(self) -> Iterator[pd.DataFrame]:
@@ -255,80 +259,94 @@ def _reset_batch_cache(self) -> None:
 
     def _reset_batches_for_new_page_size(self) -> None:
         """Reset the batch iterator when page size changes."""
-        self._batches = self._dataframe._to_pandas_batches(page_size=self.page_size)
+        self._batches = self._dataframe.to_pandas_batches(page_size=self.page_size)
 
         self._reset_batch_cache()
 
     def _set_table_html(self) -> None:
         """Sets the current html data based on the current page and page size."""
-        if self._error_message:
-            self.table_html = (
-                f"<div class='bigframes-error-message'>{self._error_message}</div>"
-            )
+        if hasattr(self, "_setting_html"):
             return
+        self._setting_html = True
 
-        # Apply sorting if a column is selected
-        df_to_display = self._dataframe
-        if self.sort_column:
-            # TODO(b/463715504): Support sorting by index columns.
-            df_to_display = df_to_display.sort_values(
-                by=self.sort_column, ascending=self.sort_ascending
-            )
-
-        # Reset batches when sorting changes
-        if self._last_sort_state != _SortState(self.sort_column, self.sort_ascending):
-            self._batches = df_to_display._to_pandas_batches(page_size=self.page_size)
-            self._reset_batch_cache()
-            self._last_sort_state = _SortState(self.sort_column, self.sort_ascending)
-            self.page = 0  # Reset to first page
-
-        start = self.page * self.page_size
-        end = start + self.page_size
-
-        # fetch more data if the requested page is outside our cache
-        cached_data = self._cached_data
-        while len(cached_data) < end and not self._all_data_loaded:
-            if self._get_next_batch():
-                cached_data = self._cached_data
+        try:
+            if self._error_message:
+                self.table_html = (
+                    f"<div class='bigframes-error-message'>{self._error_message}</div>"
+                )
+                return
+
+            # Apply sorting if a column is selected
+            df_to_display = self._dataframe
+            if self.sort_column:
+                # TODO(b/463715504): Support sorting by index columns.
+                df_to_display = df_to_display.sort_values(
+                    by=self.sort_column, ascending=self.sort_ascending
+                )
+
+            # Reset batches when sorting changes
+            if self._last_sort_state != _SortState(
+                self.sort_column, self.sort_ascending
+            ):
+                self._batches = df_to_display.to_pandas_batches(
+                    page_size=self.page_size
+                )
+                self._reset_batch_cache()
+                self._last_sort_state = _SortState(
+                    self.sort_column, self.sort_ascending
+                )
+                self.page = 0  # Reset to first page
+
+            start = self.page * self.page_size
+            end = start + self.page_size
+
+            # fetch more data if the requested page is outside our cache
+            cached_data = self._cached_data
+            while len(cached_data) < end and not self._all_data_loaded:
+                if self._get_next_batch():
+                    cached_data = self._cached_data
+                else:
+                    break
+
+            # Get the data for the current page
+            page_data = cached_data.iloc[start:end].copy()
+
+            # Handle index display
+            # TODO(b/438181139): Add tests for custom multiindex
+            if self._dataframe._block.has_index:
+                index_name = page_data.index.name
+                page_data.insert(
+                    0, index_name if index_name is not None else "", page_data.index
+                )
             else:
-                break
-
-        # Get the data for the current page
-        page_data = cached_data.iloc[start:end].copy()
-
-        # Handle index display
-        # TODO(b/438181139): Add tests for custom multiindex
-        if self._dataframe._block.has_index:
-            index_name = page_data.index.name
-            page_data.insert(
-                0, index_name if index_name is not None else "", page_data.index
+                # Default index - include as "Row" column
+                page_data.insert(0, "Row", range(start + 1, start + len(page_data) + 1))
+            # Handle case where user navigated beyond available data with unknown row count
+            is_unknown_count = self.row_count is None
+            is_beyond_data = (
+                self._all_data_loaded and len(page_data) == 0 and self.page > 0
             )
-        else:
-            # Default index - include as "Row" column
-            page_data.insert(0, "Row", range(start + 1, start + len(page_data) + 1))
-        # Handle case where user navigated beyond available data with unknown row count
-        is_unknown_count = self.row_count is None
-        is_beyond_data = self._all_data_loaded and len(page_data) == 0 and self.page > 0
-        if is_unknown_count and is_beyond_data:
-            # Calculate the last valid page (zero-indexed)
-            total_rows = len(cached_data)
-            if total_rows > 0:
-                last_valid_page = max(0, math.ceil(total_rows / self.page_size) - 1)
-                # Navigate back to the last valid page
-                self.page = last_valid_page
-                # Recursively call to display the correct page
-                return self._set_table_html()
-            else:
-                # If no data at all, stay on page 0 with empty display
-                self.page = 0
-                return self._set_table_html()
-
-        # Generate HTML table
-        self.table_html = bigframes.display.html.render_html(
-            dataframe=page_data,
-            table_id=f"table-{self._table_id}",
-            orderable_columns=self.orderable_columns,
-        )
+            if is_unknown_count and is_beyond_data:
+                # Calculate the last valid page (zero-indexed)
+                total_rows = len(cached_data)
+                if total_rows > 0:
+                    last_valid_page = max(0, math.ceil(total_rows / self.page_size) - 1)
+                    # Navigate back to the last valid page
+                    self.page = last_valid_page
+                    # Recursively call to display the correct page
+                    return self._set_table_html()
+                else:
+                    # If no data at all, stay on page 0 with empty display
+                    self.page = 0
+                    return self._set_table_html()
+
+            # Generate HTML table
+            self.table_html = bigframes.display.html.render_html(
+                dataframe=page_data,
+                table_id=f"table-{self._table_id}",
+            )
+        finally:
+            delattr(self, "_setting_html")
 
     @traitlets.observe("sort_column", "sort_ascending")
     def _sort_changed(self, _change: Dict[str, Any]):