change return type

Author: shuoweil

bigframes/core/blocks.py

@@ -99,7 +99,7 @@
 
 
 @dataclasses.dataclass
-class PandasBatches(Iterator[pd.DataFrame]):
+class PandasBatches:
     """Interface for mutable objects with state represented by a block value object."""
 
     def __init__(
@@ -124,6 +124,9 @@ def total_bytes_processed(self) -> Optional[int]:
     def __next__(self) -> pd.DataFrame:
         return next(self._dataframes)
 
+    def __iter__(self) -> Iterator[pd.DataFrame]:
+        return self
+
 
 @dataclasses.dataclass()
 class MaterializationOptions:
@@ -693,7 +696,7 @@ def to_pandas_batches(
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
         allow_large_results: Optional[bool] = None,
-    ) -> Iterator[pd.DataFrame]:
+    ) -> PandasBatches:
         """Download results one message at a time.
 
         page_size and max_results determine the size and number of batches,

bigframes/dataframe.py

@@ -1892,7 +1892,7 @@ def to_pandas_batches(
         max_results: Optional[int] = None,
         *,
         allow_large_results: Optional[bool] = None,
-    ) -> Iterable[pandas.DataFrame]:
+    ) -> blocks.PandasBatches:
         """Stream DataFrame results to an iterable of pandas DataFrame.
 
         page_size and max_results determine the size and number of batches,
@@ -1937,7 +1937,7 @@ def to_pandas_batches(
                 over the default size limit of 10 GB.
 
         Returns:
-            Iterable[pandas.DataFrame]:
+            bigframes.core.blocks.PandasBatches:
                 An iterable of smaller dataframes which combine to
                 form the original dataframe. Results stream from bigquery,
                 see https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.table.RowIterator#google_cloud_bigquery_table_RowIterator_to_arrow_iterable

bigframes/display/anywidget.py

@@ -17,15 +17,13 @@
 from importlib import resources
 import functools
 import math
-from typing import Any, cast, Dict, Iterator, List, Optional, Type
+from typing import Any, Dict, Iterator, List, Optional, Type
 import uuid
 
 import pandas as pd
 
 import bigframes
-import bigframes.core.blocks
 import bigframes.display.html
-import bigframes.session.execution_spec
 
 # 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
@@ -90,21 +88,11 @@ def __init__(self, dataframe: bigframes.dataframe.DataFrame):
         # Respect display options for initial page size
         self.page_size = bigframes.options.display.max_rows
 
-        # Force execution with explicit destination to get total_rows metadata
-        execute_result = dataframe._block.session._executor.execute(
-            dataframe._block.expr,
-            execution_spec=bigframes.session.execution_spec.ExecutionSpec(
-                ordered=True, promise_under_10gb=False
-            ),
-        )
         # The query issued by `to_pandas_batches()` already contains
         # metadata about how many results there were. Use that to avoid
         # doing an extra COUNT(*) query that `len(...)` would do.
-        self.row_count = execute_result.total_rows or 0
-        self._batches = cast(
-            bigframes.core.blocks.PandasBatches,
-            execute_result.to_pandas_batches(page_size=self.page_size),
-        )
+        self._batches = self._dataframe.to_pandas_batches(page_size=self.page_size)
+        self.row_count = self._batches.total_rows or 0
 
         self._set_table_html()
         self._initial_load_complete = True
@@ -198,19 +186,7 @@ def _cached_data(self) -> pd.DataFrame:
 
     def _reset_batches_for_new_page_size(self) -> None:
         """Reset the batch iterator when page size changes."""
-        # Execute with explicit destination for consistency with __init__
-        execute_result = self._dataframe._block.session._executor.execute(
-            self._dataframe._block.expr,
-            execution_spec=bigframes.session.execution_spec.ExecutionSpec(
-                ordered=True, promise_under_10gb=False
-            ),
-        )
-
-        # Create pandas batches from the ExecuteResult
-        self._batches = cast(
-            bigframes.core.blocks.PandasBatches,
-            execute_result.to_pandas_batches(page_size=self.page_size),
-        )
+        self._batches = self._dataframe.to_pandas_batches(page_size=self.page_size)
 
         self._cached_batches = []
         self._batch_iter = None

notebooks/dataframes/dataframe.ipynb

@@ -5366,7 +5366,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "3.10.18",
    "language": "python",
    "name": "python3"
   },
@@ -5380,7 +5380,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.1"
+   "version": "3.10.18"
   }
  },
  "nbformat": 4,