fix: resolve deferred mode display & thread execution reviews

Author: shuoweil

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
@@ -157,25 +160,31 @@ def _on_start_execution(self, change: dict[str, Any]):
 
             def run_execution():
                 try:
+                    self._error_message = None
                     if self.is_deferred_mode:
-                        self.is_deferred_mode = False
                         if self._deferred_dataframe is not None:
                             result = self._deferred_dataframe.execute()
                             if isinstance(result, bigframes.series.Series):
                                 df = result.to_frame()
-                            else:
+                            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.is_deferred_mode = False
                     self.start_execution = False
 
             self._execution_thread = threading.Thread(target=run_execution, daemon=True)

packages/bigframes/bigframes/display/html.py

@@ -245,8 +245,14 @@ def get_anywidget_bundle(
     else:
         df = obj
 
-    widget = display.TableWidget(df, dry_run_info=dry_run_info)
-    display_df, _ = df._process_display_df()
+    from bigframes.session import deferred
+
+    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):

packages/bigframes/bigframes/display/table_widget_angular/src/app/app.ts

@@ -495,18 +495,20 @@ export class App {
       const _html = this.state.tableHtml();
       const _sort = this.state.sortContext();
       const _orderable = this.state.orderableColumns();
-      const deferred = this.isDeferredMode();
-
-      if (deferred) {
-        this.isHeightInitialized = false;
-      }
+      this.isHeightInitialized = false;
 
       // Schedule DOM post-processing once the innerHTML render completes
       setTimeout(() => {
         this.applySortIndicators();
         this.lockInitialHeight();
       }, 0);
     });
+
+    effect(() => {
+      if (!this.state.startExecution()) {
+        this.isLoading.set(false);
+      }
+    });
   }
 
   ngOnInit() {

packages/bigframes/mypy.ini

@@ -44,3 +44,6 @@ ignore_missing_imports = True
 
 [mypy-anywidget]
 ignore_missing_imports = True
+
+[mypy-bigframes_vendored.*]
+ignore_errors = True

packages/bigframes/tests/unit/display/test_anywidget.py

@@ -286,7 +286,7 @@ def __init__(self):
                 pass
 
         df = mock.MagicMock(spec=MockDeferredBigQueryDataFrame)
-        df.__class__ = DeferredBigQueryDataFrame
+        df.__class__ = DeferredBigQueryDataFrame  # type: ignore[assignment]
         yield df
 
 
@@ -458,5 +458,5 @@ def test_deferred_mode_execution_error(mock_deferred_df):
         if thread is not None:
             thread.join(timeout=5)
 
-        assert widget.is_deferred_mode is False
+        assert widget.is_deferred_mode is True
         assert widget._error_message == "Query Failed"