feat: support interactive execution of deferred DataFrames in TableWidget

packages/bigframes/bigframes/dataframe.py

@@ -819,7 +819,7 @@ def __repr__(self) -> str:
             column_count=len(self.columns),
         )
 
-    def _prepare_display_df(self) -> DataFrame:
+    def _process_display_df(self) -> tuple[DataFrame, list[str]]:
         """Process ObjectRef and JSON/nested JSON columns for display."""
         df = self
         # Arrow/Pandas to_pandas_batches does not support raw JSON/nested JSON
@@ -837,7 +837,7 @@ def _prepare_display_df(self) -> DataFrame:
                 sql_template="TO_JSON_STRING({0})",
             )
             df = df.assign(**{col: df[col]._apply_unary_op(op) for col in json_cols})
-        return df
+        return df, []
 
     def _repr_mimebundle_(self, include=None, exclude=None):
         """

packages/bigframes/bigframes/display/anywidget.py

@@ -18,6 +18,9 @@
 
 import dataclasses
 import functools
+import logging
+
+logger = logging.getLogger(__name__)
 import math
 import threading
 import uuid
@@ -77,8 +80,18 @@ class TableWidget(_WIDGET_BASE):
     _error_message = traitlets.Unicode(allow_none=True, default_value=None).tag(
         sync=True
     )
-
-    def __init__(self, dataframe: bigframes.dataframe.DataFrame):
+    start_execution = traitlets.Bool(False).tag(sync=True)
+    is_deferred_mode = traitlets.Bool(False).tag(sync=True)
+    dry_run_info = traitlets.Unicode("").tag(sync=True)
+
+    def __init__(
+        self,
+        dataframe: (
+            bigframes.dataframe.DataFrame
+            | bigframes.session.deferred.DeferredBigQueryDataFrame
+        ),
+        dry_run_info: Optional[str] = None,
+    ):
         """Initialize the TableWidget.
 
         Args:
@@ -90,14 +103,34 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
                 "`pip install 'bigframes[anywidget]'` to use TableWidget."
             )
 
-        self._dataframe = dataframe
+        from bigframes.session import deferred
+
+        is_deferred = False
+        deferred_df = None
+        df = None
+
+        if isinstance(dataframe, deferred.DeferredBigQueryDataFrame):
+            is_deferred = True
+            deferred_df = dataframe
+        elif bigframes.options.display.repr_mode == "deferred":
+            is_deferred = True
+            df = dataframe
+        else:
+            df = dataframe
 
         from bigframes.core.utils import get_ipython_execution_count
 
         self._cell_execution_count = get_ipython_execution_count()
 
         super().__init__()
 
+        self.is_deferred_mode = is_deferred
+        self._deferred_dataframe = deferred_df
+        self._dataframe = df
+
+        if dry_run_info:
+            self.dry_run_info = dry_run_info
+
         # Initialize attributes that might be needed by observers first
         self._table_id = str(uuid.uuid4())
         self._all_data_loaded = False
@@ -111,19 +144,60 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         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
 
-        self.orderable_columns = self._get_orderable_columns(dataframe)
-
-        self._initial_load()
+        if not self.is_deferred_mode:
+            self._initialize_from_dataframe()
 
         # Signals to the frontend that the initial data load is complete.
         # Also used as a guard to prevent observers from firing during initialization.
         self._initial_load_complete = True
 
+    @traitlets.observe("start_execution")
+    def _on_start_execution(self, change: dict[str, Any]):
+        if change["new"]:
+
+            def run_execution():
+                try:
+                    self._error_message = None
+                    if self.is_deferred_mode:
+                        if self._deferred_dataframe is not None:
+                            result = self._deferred_dataframe.execute()
+                            if isinstance(result, bigframes.series.Series):
+                                df = result.to_frame()
+                            elif isinstance(result, bigframes.dataframe.DataFrame):
+                                df = result
+                            else:
+                                raise TypeError(
+                                    f"Unexpected result type: {type(result)}"
+                                )
+                            self._dataframe, _ = df._process_display_df()
+                            self._initialize_from_dataframe()
+                            self.is_deferred_mode = False
+                        elif self._dataframe is not None:
+                            self._dataframe, _ = self._dataframe._process_display_df()
+                            self._initialize_from_dataframe()
+                            self.is_deferred_mode = False
+                    elif not self.is_deferred_mode and self._dataframe is not None:
+                        self._initial_load()
+                except Exception as e:
+                    logger.warning(f"Error in background execution: {e}")
+                    self._error_message = str(e)
+                finally:
+                    self.start_execution = False
+
+            self._execution_thread = threading.Thread(target=run_execution, daemon=True)
+            self._execution_thread.start()
+
+    def _initialize_from_dataframe(self):
+        if self._dataframe is None:
+            return
+
+        self.orderable_columns = self._get_orderable_columns(self._dataframe)
+
+        self._initial_load()
+
     def _get_orderable_columns(
         self, dataframe: bigframes.dataframe.DataFrame
     ) -> list[str]:
@@ -278,7 +352,9 @@ def _batch_iterator(self) -> Iterator[pd.DataFrame]:
     def _cached_data(self) -> pd.DataFrame:
         """Combine all cached batches into a single DataFrame."""
         if not self._cached_batches:
-            return pd.DataFrame(columns=self._dataframe.columns)
+            if self._dataframe is not None:
+                return pd.DataFrame(columns=self._dataframe.columns)
+            return pd.DataFrame()
         return pd.concat(self._cached_batches)
 
     def _reset_batch_cache(self) -> None:
@@ -289,6 +365,8 @@ def _reset_batch_cache(self) -> None:
 
     def _reset_batches_for_new_page_size(self) -> None:
         """Reset the batch iterator when page size changes."""
+        if self._dataframe is None:
+            return
         with bigframes.option_context("display.progress_bar", None):
             self._batches = self._dataframe.to_pandas_batches(
                 page_size=self.page_size,
@@ -299,6 +377,9 @@ def _reset_batches_for_new_page_size(self) -> None:
 
     def _set_table_html(self) -> None:
         """Sets the current html data based on the current page and page size."""
+        if self.is_deferred_mode:
+            return
+
         new_page = None
         with (
             self._setting_html_lock,
@@ -310,6 +391,10 @@ def _set_table_html(self) -> None:
                 )
                 return
 
+            if self._dataframe is None:
+                self.table_html = "<div class='bigframes-error-message'>Internal Error: DataFrame is missing.</div>"
+                return
+
             # Apply sorting if a column is selected
             df_to_display = self._dataframe
             sort_columns = [item["column"] for item in self.sort_context]

packages/bigframes/bigframes/display/html.py

@@ -231,6 +231,7 @@ def get_anywidget_bundle(
     obj: Union[bigframes.dataframe.DataFrame, bigframes.series.Series],
     include=None,
     exclude=None,
+    dry_run_info: str | None = None,
 ) -> tuple[dict[str, Any], dict[str, Any]]:
     """
     Helper method to create and return the anywidget mimebundle.
@@ -244,9 +245,17 @@ def get_anywidget_bundle(
     else:
         df = obj
 
-    df = df._prepare_display_df()
+    from bigframes.session import deferred
 
-    widget = display.TableWidget(df)
+    if (
+        not isinstance(df, deferred.DeferredBigQueryDataFrame)
+        and bigframes.options.display.repr_mode != "deferred"
+    ):
+        display_df, _ = df._process_display_df()
+    else:
+        display_df = df
+
+    widget = display.TableWidget(display_df, dry_run_info=dry_run_info)
     widget_repr_result = widget._repr_mimebundle_(include=include, exclude=exclude)
 
     if isinstance(widget_repr_result, tuple):
@@ -262,20 +271,23 @@ def get_anywidget_bundle(
     total_rows = widget.row_count
     total_columns = len(df.columns)
 
-    widget_repr["text/html"] = create_html_representation(
-        obj,
-        cached_pd,
-        total_rows,
-        total_columns,
-    )
-    is_series, has_index = _get_obj_metadata(obj)
-    widget_repr["text/plain"] = plaintext.create_text_representation(
-        cached_pd,
-        total_rows,
-        is_series=is_series,
-        has_index=has_index,
-        column_count=len(df.columns) if not is_series else 0,
-    )
+    if dry_run_info:
+        widget_repr["text/plain"] = dry_run_info
+    else:
+        widget_repr["text/html"] = create_html_representation(
+            obj,
+            cached_pd,
+            total_rows,
+            total_columns,
+        )
+        is_series, has_index = _get_obj_metadata(obj)
+        widget_repr["text/plain"] = plaintext.create_text_representation(
+            cached_pd,
+            total_rows,
+            is_series=is_series,
+            has_index=has_index,
+            column_count=len(df.columns) if not is_series else 0,
+        )
 
     return widget_repr, widget_metadata
 
@@ -300,7 +312,7 @@ def repr_mimebundle_head(
     else:
         df = obj
 
-    df = df._prepare_display_df()
+    df, _ = df._process_display_df()
     pandas_df, row_count, query_job = df._block.retrieve_repr_request_results(
         opts.max_rows
     )
@@ -332,27 +344,40 @@ def repr_mimebundle(
     # BQ Studio, but there is a known compatibility issue with Marimo that needs to be addressed.
 
     opts = options.display
-    if opts.repr_mode == "deferred":
-        return repr_mimebundle_deferred(obj)
-
-    if opts.render_mode == "anywidget" or opts.repr_mode == "anywidget":
+    if (
+        opts.render_mode == "anywidget"
+        or opts.repr_mode == "anywidget"
+        or opts.repr_mode == "deferred"
+    ):
         try:
             with bigframes.option_context("display.progress_bar", None):
                 with warnings.catch_warnings():
                     warnings.simplefilter(
                         "ignore", category=bigframes.exceptions.JSONDtypeWarning
                     )
                     warnings.simplefilter("ignore", category=FutureWarning)
-                    return get_anywidget_bundle(obj, include=include, exclude=exclude)
-        except ImportError:
+                    dry_run_info = None
+                    if opts.repr_mode == "deferred":
+                        dry_run_job = obj._compute_dry_run()
+                        dry_run_info = formatter.repr_query_job(dry_run_job)
+                    return get_anywidget_bundle(
+                        obj,
+                        include=include,
+                        exclude=exclude,
+                        dry_run_info=dry_run_info,
+                    )
+        except Exception:
             # 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. "
+                "Anywidget mode is not available or failed to load. "
+                "Please `pip install anywidget traitlets` or "
+                "`pip install 'bigframes[anywidget]'` to use interactive tables. "
                 f"Falling back to static HTML. Error: {traceback.format_exc()}"
             )
+            if opts.repr_mode == "deferred":
+                return repr_mimebundle_deferred(obj)
 
     bundle = repr_mimebundle_head(obj)
     if opts.render_mode == "plaintext":