fix: Correct DataFrame widget rendering in Colab (#2319)

This PR fix the _get_anywidget_bundle method. Previously, when the
underlying widget's _repr_mimebundle_ method returned a (data, metadata)
tuple, the code was only extracting the data portion and discarding the
metadata. This resulted in the widget not rendering correctly in
environments like Colab, which rely on this metadata.

The change corrects this by properly unpacking the tuple into
widget_repr and widget_metadata. The method now preserves the metadata
and returns it along with the data, ensuring that the necessary
information for widget rendering is passed on.

We also revert commit 4df3428 to
reapply "refactor: Migrate DataFrame display to use IPython's
repr_mimebundle() protocol for anywidget mode
(#2271)"


A testcase is added to verify this new change. We also verified at
colab: screen/AzGa5RMTJnMH5NH

Fixes #<466155761> 🦕

Author: shuoweil

bigframes/core/indexes/base.py

@@ -376,9 +376,7 @@ def __repr__(self) -> __builtins__.str:
         # metadata, like we do with DataFrame.
         opts = bigframes.options.display
         max_results = opts.max_rows
-        # anywdiget mode uses the same display logic as the "deferred" mode
-        # for faster execution
-        if opts.repr_mode in ("deferred", "anywidget"):
+        if opts.repr_mode == "deferred":
             _, dry_run_query_job = self._block._compute_dry_run()
             return formatter.repr_query_job(dry_run_query_job)
 

bigframes/dataframe.py

@@ -789,9 +789,7 @@ def __repr__(self) -> str:
 
         opts = bigframes.options.display
         max_results = opts.max_rows
-        # anywdiget mode uses the same display logic as the "deferred" mode
-        # for faster execution
-        if opts.repr_mode in ("deferred", "anywidget"):
+        if opts.repr_mode == "deferred":
             return formatter.repr_query_job(self._compute_dry_run())
 
         # TODO(swast): pass max_columns and get the true column count back. Maybe
@@ -829,68 +827,149 @@ def __repr__(self) -> str:
         lines.append(f"[{row_count} rows x {column_count} columns]")
         return "\n".join(lines)
 
-    def _repr_html_(self) -> str:
-        """
-        Returns an html string primarily for use by notebooks for displaying
-        a representation of the DataFrame. Displays 20 rows by default since
-        many notebooks are not configured for large tables.
-        """
-        opts = bigframes.options.display
-        max_results = opts.max_rows
-        if opts.repr_mode == "deferred":
-            return formatter.repr_query_job(self._compute_dry_run())
-
-        # Process blob columns first, regardless of display mode
-        self._cached()
-        df = self.copy()
+    def _get_display_df_and_blob_cols(self) -> tuple[DataFrame, list[str]]:
+        """Process blob columns for display."""
+        df = self
+        blob_cols = []
         if bigframes.options.display.blob_display:
             blob_cols = [
                 series_name
-                for series_name, series in df.items()
+                for series_name, series in self.items()
                 if series.dtype == bigframes.dtypes.OBJ_REF_DTYPE
             ]
-            for col in blob_cols:
-                # TODO(garrettwu): Not necessary to get access urls for all the rows. Update when having a to get URLs from local data.
-                df[col] = df[col].blob._get_runtime(mode="R", with_metadata=True)
+            if blob_cols:
+                df = self.copy()
+                for col in blob_cols:
+                    # TODO(garrettwu): Not necessary to get access urls for all the rows. Update when having a to get URLs from local data.
+                    df[col] = df[col].blob._get_runtime(mode="R", with_metadata=True)
+        return df, blob_cols
+
+    def _get_anywidget_bundle(
+        self, include=None, exclude=None
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """
+        Helper method to create and return the anywidget mimebundle.
+        This function encapsulates the logic for anywidget display.
+        """
+        from bigframes import display
+
+        df, blob_cols = self._get_display_df_and_blob_cols()
+
+        # Create and display the widget
+        widget = display.TableWidget(df)
+        widget_repr_result = widget._repr_mimebundle_(include=include, exclude=exclude)
+
+        # Handle both tuple (data, metadata) and dict returns
+        if isinstance(widget_repr_result, tuple):
+            widget_repr, widget_metadata = widget_repr_result
         else:
-            blob_cols = []
+            widget_repr = widget_repr_result
+            widget_metadata = {}
+
+        widget_repr = dict(widget_repr)
+
+        # At this point, we have already executed the query as part of the
+        # widget construction. Let's use the information available to render
+        # the HTML and plain text versions.
+        widget_repr["text/html"] = self._create_html_representation(
+            widget._cached_data,
+            widget.row_count,
+            len(self.columns),
+            blob_cols,
+        )
 
-        if opts.repr_mode == "anywidget":
-            try:
-                from IPython.display import display as ipython_display
+        widget_repr["text/plain"] = self._create_text_representation(
+            widget._cached_data, widget.row_count
+        )
+
+        return widget_repr, widget_metadata
+
+    def _create_text_representation(
+        self, pandas_df: pandas.DataFrame, total_rows: typing.Optional[int]
+    ) -> str:
+        """Create a text representation of the DataFrame."""
+        opts = bigframes.options.display
+        with display_options.pandas_repr(opts):
+            import pandas.io.formats
+
+            # safe to mutate this, this dict is owned by this code, and does not affect global config
+            to_string_kwargs = (
+                pandas.io.formats.format.get_dataframe_repr_params()  # type: ignore
+            )
+            if not self._has_index:
+                to_string_kwargs.update({"index": False})
+
+            # We add our own dimensions string, so don't want pandas to.
+            to_string_kwargs.update({"show_dimensions": False})
+            repr_string = pandas_df.to_string(**to_string_kwargs)
 
-                from bigframes import display
+        lines = repr_string.split("\n")
 
-                # Always create a new widget instance for each display call
-                # This ensures that each cell gets its own widget and prevents
-                # unintended sharing between cells
-                widget = display.TableWidget(df.copy())
+        if total_rows is not None and total_rows > len(pandas_df):
+            lines.append("...")
 
-                ipython_display(widget)
-                return ""  # Return empty string since we used display()
+        lines.append("")
+        column_count = len(self.columns)
+        lines.append(f"[{total_rows or '?'} rows x {column_count} columns]")
+        return "\n".join(lines)
 
-            except (AttributeError, ValueError, ImportError):
-                # Fallback if anywidget is not available
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        """
+        Custom display method for IPython/Jupyter environments.
+        This is called by IPython's display system when the object is displayed.
+        """
+        # TODO(b/467647693): Anywidget integration has been tested in Jupyter, VS Code, and
+        # BQ Studio, but there is a known compatibility issue with Marimo that needs to be addressed.
+        opts = bigframes.options.display
+        # Only handle widget display in anywidget mode
+        if opts.repr_mode == "anywidget":
+            try:
+                return self._get_anywidget_bundle(include=include, exclude=exclude)
+
+            except ImportError:
+                # Anywidget is an optional dependency, so warn rather than fail.
+                # TODO(shuowei): When Anywidget becomes the default for all repr modes,
+                # remove this warning.
                 warnings.warn(
                     "Anywidget mode is not available. "
                     "Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'` to use interactive tables. "
-                    f"Falling back to deferred mode. Error: {traceback.format_exc()}"
+                    f"Falling back to static HTML. Error: {traceback.format_exc()}"
                 )
-                return formatter.repr_query_job(self._compute_dry_run())
 
-        # Continue with regular HTML rendering for non-anywidget modes
-        # TODO(swast): pass max_columns and get the true column count back. Maybe
-        # get 1 more column than we have requested so that pandas can add the
-        # ... for us?
+        # In non-anywidget mode, fetch data once and use it for both HTML
+        # and plain text representations to avoid multiple queries.
+        opts = bigframes.options.display
+        max_results = opts.max_rows
+
+        df, blob_cols = self._get_display_df_and_blob_cols()
+
         pandas_df, row_count, query_job = df._block.retrieve_repr_request_results(
             max_results
         )
-
         self._set_internal_query_job(query_job)
         column_count = len(pandas_df.columns)
 
+        html_string = self._create_html_representation(
+            pandas_df, row_count, column_count, blob_cols
+        )
+
+        text_representation = self._create_text_representation(pandas_df, row_count)
+
+        return {"text/html": html_string, "text/plain": text_representation}
+
+    def _create_html_representation(
+        self,
+        pandas_df: pandas.DataFrame,
+        row_count: int,
+        column_count: int,
+        blob_cols: list[str],
+    ) -> str:
+        """Create an HTML representation of the DataFrame."""
+        opts = bigframes.options.display
         with display_options.pandas_repr(opts):
-            # Allows to preview images in the DataFrame. The implementation changes the string repr as well, that it doesn't truncate strings or escape html charaters such as "<" and ">". We may need to implement a full-fledged repr module to better support types not in pandas.
+            # TODO(shuowei, b/464053870): Escaping HTML would be useful, but
+            # `escape=False` is needed to show images. We may need to implement
+            # a full-fledged repr module to better support types not in pandas.
             if bigframes.options.display.blob_display and blob_cols:
 
                 def obj_ref_rt_to_html(obj_ref_rt) -> str:
@@ -919,15 +998,12 @@ def obj_ref_rt_to_html(obj_ref_rt) -> str:
 
                 # set max_colwidth so not to truncate the image url
                 with pandas.option_context("display.max_colwidth", None):
-                    max_rows = pandas.get_option("display.max_rows")
-                    max_cols = pandas.get_option("display.max_columns")
-                    show_dimensions = pandas.get_option("display.show_dimensions")
                     html_string = pandas_df.to_html(
                         escape=False,
                         notebook=True,
-                        max_rows=max_rows,
-                        max_cols=max_cols,
-                        show_dimensions=show_dimensions,
+                        max_rows=pandas.get_option("display.max_rows"),
+                        max_cols=pandas.get_option("display.max_columns"),
+                        show_dimensions=pandas.get_option("display.show_dimensions"),
                         formatters=formatters,  # type: ignore
                     )
             else:

bigframes/streaming/dataframe.py

@@ -291,13 +291,13 @@ def __repr__(self, *args, **kwargs):
 
     __repr__.__doc__ = _curate_df_doc(inspect.getdoc(dataframe.DataFrame.__repr__))
 
-    def _repr_html_(self, *args, **kwargs):
-        return _return_type_wrapper(self._df._repr_html_, StreamingDataFrame)(
+    def _repr_mimebundle_(self, *args, **kwargs):
+        return _return_type_wrapper(self._df._repr_mimebundle_, StreamingDataFrame)(
             *args, **kwargs
         )
 
-    _repr_html_.__doc__ = _curate_df_doc(
-        inspect.getdoc(dataframe.DataFrame._repr_html_)
+    _repr_mimebundle_.__doc__ = _curate_df_doc(
+        inspect.getdoc(dataframe.DataFrame._repr_mimebundle_)
     )
 
     @property