update the method of using PandasBatches.total_rows

Author: shuoweil

bigframes/display/anywidget.py

@@ -17,12 +17,14 @@
 from importlib import resources
 import functools
 import math
-from typing import Any, Dict, Iterator, List, Optional, Type
+import typing
+from typing import Any, cast, Dict, Iterator, List, Optional, Type
 import uuid
 
 import pandas as pd
 
 import bigframes
+import bigframes.core.blocks
 
 # anywidget and traitlets are optional dependencies. We don't want the import of this
 # module to fail if they aren't installed, though. Instead, we try to limit the surface that
@@ -68,21 +70,24 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         self._table_id = str(uuid.uuid4())
         self._all_data_loaded = False
         self._batch_iter: Optional[Iterator[pd.DataFrame]] = None
+        self._batches: Optional[bigframes.core.blocks.PandasBatches] = None
         self._cached_batches: List[pd.DataFrame] = []
 
         # Respect display options for initial page size
         initial_page_size = bigframes.options.display.max_rows
 
         try:
             # Fetches initial data batches and row count for display.
-            # `to_pandas_batches` provides an iterable of pandas DataFrames
-            # and eagerly retrieves the total row count
-            self._batches = dataframe.to_pandas_batches(
+            batches = dataframe.to_pandas_batches(
                 page_size=initial_page_size,
             )
+            self._batches = cast(bigframes.core.blocks.PandasBatches, batches)
 
-            # Access the total_rows property directly
-            self.row_count = self._batches.total_rows or 0
+            # Use total_rows if available, otherwise default to 0.
+            if self._batches:
+                self.row_count = self._batches.total_rows or 0
+            else:
+                self.row_count = 0
             self.page_size = initial_page_size
 
             # Generates the initial HTML table content
@@ -91,7 +96,7 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         except Exception:
             self.row_count = 0
             self.page_size = initial_page_size
-            self._batches = iter([])
+            self._batches = None
             self.table_html = ""
 
     @functools.cached_property
@@ -175,7 +180,10 @@ def _get_next_batch(self) -> bool:
     def _batch_iterator(self) -> Iterator[pd.DataFrame]:
         """Lazily initializes and returns the batch iterator."""
         if self._batch_iter is None:
-            self._batch_iter = iter(self._batches)
+            if self._batches is None:
+                self._batch_iter = iter([])
+            else:
+                self._batch_iter = iter(self._batches)
         return self._batch_iter
 
     @property
@@ -187,7 +195,8 @@ def _cached_data(self) -> pd.DataFrame:
 
     def _reset_batches_for_new_page_size(self):
         """Reset the batch iterator when page size changes."""
-        self._batches = self._dataframe.to_pandas_batches(page_size=self.page_size)
+        batches = self._dataframe.to_pandas_batches(page_size=self.page_size)
+        self._batches = typing.cast(bigframes.core.blocks.PandasBatches, batches)
         self._cached_batches = []
         self._batch_iter = None
         self._all_data_loaded = False

tests/system/small/test_anywidget.py

@@ -476,14 +476,21 @@ def test_widget_row_count_should_be_immutable_after_creation(
     assert widget.row_count == initial_row_count
 
 
+class FaultyIterator:
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        raise ValueError("Simulated read error")
+
+
 @pytest.mark.parametrize(
     "total_rows_param, arrow_batches_param",
     [
-        # Corresponds to mock_execute_total_rows_is_none
+        # Case 1: total_rows is None, which should be handled gracefully.
         (None, []),
-        # Corresponds to mock_execute_batches_are_invalid (assuming empty list
-        # for invalid batches for now)
-        (100, []),
+        # Case 2: Batches are invalid and will raise an error during iteration.
+        (100, FaultyIterator()),
     ],
     ids=[
         "when_total_rows_is_None",