change tests in read_gbq_colab

Author: shuoweil

bigframes/display/anywidget.py

@@ -77,28 +77,21 @@ def __init__(self, dataframe: bigframes.dataframe.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.
-            batches = dataframe.to_pandas_batches(
-                page_size=initial_page_size,
-            )
-            self._batches = cast(bigframes.core.blocks.PandasBatches, batches)
+        # Fetches initial data batches and row count for display.
+        batches = dataframe.to_pandas_batches(
+            page_size=initial_page_size,
+        )
+        self._batches = cast(bigframes.core.blocks.PandasBatches, batches)
 
-            # 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
+        # Use total_rwos from batches directly
+        self.row_count = self._batches.total_rows or 0
 
-            # Generates the initial HTML table content
-            self._set_table_html()
+        # Set page_size after _batches is initialized so observers have
+        # access to batch data
+        self.page_size = initial_page_size
 
-        except Exception:
-            self.row_count = 0
-            self.page_size = initial_page_size
-            self._batches = None
-            self.table_html = ""
+        # Generates the initial HTML table content
+        self._set_table_html()
 
     @functools.cached_property
     def _esm(self):

tests/benchmark/read_gbq_colab/first_page.py

@@ -27,10 +27,10 @@ def first_page(*, project_id, dataset_id, table_id):
         f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
     )
 
-    # Get number of rows (to calculate number of pages) and the first page.
-    batches = df._to_pandas_batches(page_size=PAGE_SIZE)
-    assert (tr := batches.total_rows) is not None and tr >= 0
-    next(iter(batches))
+    # Use total_rows from batches directly and the first page
+    execute_result = df._block.session._executor.execute(df._block.expr, ordered=True)
+    execute_result.total_rows or 0
+    next(iter(df.to_pandas_batches(page_size=PAGE_SIZE)))
 
 
 if __name__ == "__main__":

tests/benchmark/read_gbq_colab/last_page.py

@@ -27,10 +27,9 @@ def last_page(*, project_id, dataset_id, table_id):
         f"SELECT * FROM `{project_id}`.{dataset_id}.{table_id}"
     )
 
-    # Get number of rows (to calculate number of pages) and then all pages.
-    batches = df._to_pandas_batches(page_size=PAGE_SIZE)
-    assert (tr := batches.total_rows) is not None and tr >= 0
-    for _ in batches:
+    execute_result = df._block.session._executor.execute(df._block.expr, ordered=True)
+    execute_result.total_rows or 0
+    for _ in df.to_pandas_batches(page_size=PAGE_SIZE):
         pass