feat: support interactive execution of deferred DataFrames in TableWidget

packages/bigframes/bigframes/core/compile/sqlglot/sql/base.py

@@ -49,6 +49,13 @@
 
 def to_sql(expr: sge.Expression) -> str:
     """Generate SQL string from the given expression."""
+
+    def _flatten_null_casts(node: sge.Expression) -> sge.Expression:
+        if isinstance(node, (sge.Cast, sge.TryCast)) and str(node.to).upper() == "NULL":
+            return sge.Null()
+        return node
+
+    expr = expr.transform(_flatten_null_casts)
     return expr.sql(dialect=DIALECT, pretty=PRETTY)
 
 
@@ -69,8 +76,6 @@ def literal(value: typing.Any, dtype: dtypes.Dtype | None = None) -> sge.Express
         return sge.Null()
 
     if value is None:
-        if str(sqlglot_type).upper() == "NULL":
-            return sge.Null()
         return cast(sge.Null(), sqlglot_type)
     if dtypes.is_struct_like(dtype):
         items = [
@@ -121,8 +126,10 @@ def literal(value: typing.Any, dtype: dtypes.Dtype | None = None) -> sge.Express
         return sge.convert(value)
 
 
-def cast(arg: typing.Any, to: str, safe: bool = False) -> sge.Cast | sge.TryCast:
+def cast(arg: typing.Any, to: str | sge.DataType, safe: bool = False) -> sge.Expression:
     """Return a SQL expression that casts the given argument to the specified type."""
+    if str(to).upper() == "NULL":
+        return sge.Null()
     if safe:
         return sge.TryCast(this=arg, to=to)
     else:

packages/bigframes/bigframes/dataframe.py

@@ -843,7 +843,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
@@ -861,7 +861,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,141 @@ 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"]:
+            import asyncio
+
+            try:
+                loop = asyncio.get_running_loop()
+            except RuntimeError:
+                try:
+                    import tornado.ioloop  # type: ignore[import-not-found]
+
+                    loop = tornado.ioloop.IOLoop.current().asyncio_loop  # type: ignore[attr-defined]
+                except Exception:
+                    loop = None
+
+            def run_execution():
+                try:
+                    self._error_message = None
+                    df = 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)}"
+                                )
+                        elif self._dataframe is not None:
+                            df = self._dataframe
+                    else:
+                        df = self._dataframe
+
+                    if df is None:
+                        raise ValueError("No DataFrame to execute.")
+
+                    df_to_set, _ = df._process_display_df()
+                    orderable_cols = self._get_orderable_columns(df_to_set)
+
+                    with bigframes.option_context("display.progress_bar", None):
+                        batches = df_to_set.to_pandas_batches(
+                            page_size=self.page_size,
+                            cell_execution_count=self._cell_execution_count,
+                        )
+
+                    total_rows = getattr(batches, "total_rows", None)
+
+                    # Fetch the first batch
+                    batch_iter = iter(batches)
+                    try:
+                        initial_batch = next(batch_iter)
+                        cached_batches = [initial_batch]
+                        all_data_loaded = False
+                    except StopIteration:
+                        initial_batch = pd.DataFrame(columns=df_to_set.columns)
+                        cached_batches = []
+                        all_data_loaded = True
+
+                    # Render the HTML
+                    page_data = initial_batch.copy()
+                    start = 0
+                    if df_to_set._block.has_index:
+                        is_unnamed_single_index = (
+                            page_data.index.name is None
+                            and not isinstance(page_data.index, pd.MultiIndex)
+                        )
+                        page_data = page_data.reset_index()
+                        if is_unnamed_single_index and "index" in page_data.columns:
+                            page_data.rename(columns={"index": ""}, inplace=True)
+                    else:
+                        page_data.insert(
+                            0, "Row", range(start + 1, start + len(page_data) + 1)
+                        )
+
+                    initial_html = bigframes.display.html.render_html(
+                        dataframe=page_data,
+                        table_id=f"table-{self._table_id}",
+                        orderable_columns=orderable_cols,
+                        max_columns=self.max_columns,
+                    )
+
+                    def update_ui():
+                        with self.hold_sync():
+                            self._dataframe = df_to_set
+                            self.orderable_columns = orderable_cols
+                            self._batches = batches
+                            self._batch_iter = batch_iter
+                            self._cached_batches = cached_batches
+                            self._all_data_loaded = all_data_loaded
+                            self.row_count = total_rows
+                            self.table_html = initial_html
+                            self.is_deferred_mode = False
+                            self.start_execution = False
+
+                    if loop is not None and loop.is_running():
+                        loop.call_soon_threadsafe(update_ui)
+                    else:
+                        update_ui()
+                except Exception as e:
+                    logger.warning(f"Error in background execution: {e}")
+                    err_msg = str(e)
+
+                    def set_error():
+                        with self.hold_sync():
+                            self._error_message = err_msg
+                            self.start_execution = False
+
+                    if loop is not None and loop.is_running():
+                        loop.call_soon_threadsafe(set_error)
+                    else:
+                        set_error()
+
+            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 +433,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 +446,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 +458,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 +472,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]