feat: Add total_rows property to pandas batches iterator

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -29,7 +29,17 @@
 import random
 import textwrap
 import typing
-from typing import Iterable, List, Literal, Mapping, Optional, Sequence, Tuple, Union
+from typing import (
+    Iterable,
+    Iterator,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 import warnings
 
 import bigframes_vendored.constants as constants
@@ -87,14 +97,22 @@
 LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]
 
 
-class BlockHolder(typing.Protocol):
+@dataclasses.dataclass
+class PandasBatches(Iterator[pd.DataFrame]):
     """Interface for mutable objects with state represented by a block value object."""
 
-    def _set_block(self, block: Block):
-        """Set the underlying block value of the object"""
+    def __init__(
+        self, pandas_batches: Iterator[pd.DataFrame], total_rows: Optional[int] = 0
+    ):
+        self._dataframes: Iterator[pd.DataFrame] = pandas_batches
+        self._total_rows: Optional[int] = total_rows
+
+    @property
+    def total_rows(self) -> Optional[int]:
+        return self._total_rows
 
-    def _get_block(self) -> Block:
-        """Get the underlying block value of the object"""
+    def __next__(self) -> pd.DataFrame:
+        return next(self._dataframes)
 
 
 @dataclasses.dataclass()
@@ -599,8 +617,7 @@ def try_peek(
                 self.expr, n, use_explicit_destination=allow_large_results
             )
             df = result.to_pandas()
-            self._copy_index_to_pandas(df)
-            return df
+            return self._copy_index_to_pandas(df)
         else:
             return None
 
@@ -609,8 +626,7 @@ def to_pandas_batches(
         page_size: Optional[int] = None,
         max_results: Optional[int] = None,
         allow_large_results: Optional[bool] = None,
-        squeeze: Optional[bool] = False,
-    ):
+    ) -> Iterator[pd.DataFrame]:
         """Download results one message at a time.
 
         page_size and max_results determine the size and number of batches,
@@ -621,43 +637,43 @@ def to_pandas_batches(
             use_explicit_destination=allow_large_results,
         )
 
-        total_batches = 0
-        for df in execute_result.to_pandas_batches(
-            page_size=page_size, max_results=max_results
-        ):
-            total_batches += 1
-            self._copy_index_to_pandas(df)
-            if squeeze:
-                yield df.squeeze(axis=1)
-            else:
-                yield df
-
         # To reduce the number of edge cases to consider when working with the
         # results of this, always return at least one DataFrame. See:
         # b/428918844.
-        if total_batches == 0:
-            df = pd.DataFrame(
-                {
-                    col: pd.Series([], dtype=self.expr.get_column_type(col))
-                    for col in itertools.chain(self.value_columns, self.index_columns)
-                }
-            )
-            self._copy_index_to_pandas(df)
-            yield df
+        empty_val = pd.DataFrame(
+            {
+                col: pd.Series([], dtype=self.expr.get_column_type(col))
+                for col in itertools.chain(self.value_columns, self.index_columns)
+            }
+        )
+        dfs = map(
+            lambda a: a[0],
+            itertools.zip_longest(
+                execute_result.to_pandas_batches(page_size, max_results),
+                [0],
+                fillvalue=empty_val,
+            ),
+        )
+        dfs = iter(map(self._copy_index_to_pandas, dfs))
 
-    def _copy_index_to_pandas(self, df: pd.DataFrame):
-        """Set the index on pandas DataFrame to match this block.
+        total_rows = execute_result.total_rows
+        if (total_rows is not None) and (max_results is not None):
+            total_rows = min(total_rows, max_results)
 
-        Warning: This method modifies ``df`` inplace.
-        """
+        return PandasBatches(dfs, total_rows)
+
+    def _copy_index_to_pandas(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Set the index on pandas DataFrame to match this block."""
         # Note: If BigQuery DataFrame has null index, a default one will be created for the local materialization.
+        new_df = df.copy()
         if len(self.index_columns) > 0:
-            df.set_index(list(self.index_columns), inplace=True)
+            new_df.set_index(list(self.index_columns), inplace=True)
             # Pandas names is annotated as list[str] rather than the more
             # general Sequence[Label] that BigQuery DataFrames has.
             # See: https://github.com/pandas-dev/pandas-stubs/issues/804
-            df.index.names = self.index.names  # type: ignore
-        df.columns = self.column_labels
+            new_df.index.names = self.index.names  # type: ignore
+        new_df.columns = self.column_labels
+        return new_df
 
     def _materialize_local(
         self, materialize_options: MaterializationOptions = MaterializationOptions()
@@ -724,9 +740,7 @@ def _materialize_local(
             )
         else:
             df = execute_result.to_pandas()
-            self._copy_index_to_pandas(df)
-
-        return df, execute_result.query_job
+            return self._copy_index_to_pandas(df), execute_result.query_job
 
     def _downsample(
         self, total_rows: int, sampling_method: str, fraction: float, random_state
@@ -1591,8 +1605,7 @@ def retrieve_repr_request_results(
         row_count = self.session._executor.execute(self.expr.row_count()).to_py_scalar()
 
         head_df = head_result.to_pandas()
-        self._copy_index_to_pandas(head_df)
-        return head_df, row_count, head_result.query_job
+        return self._copy_index_to_pandas(head_df), row_count, head_result.query_job
 
     def promote_offsets(self, label: Label = None) -> typing.Tuple[Block, str]:
         expr, result_id = self._expr.promote_offsets()

bigframes/series.py

@@ -648,13 +648,12 @@ def to_pandas_batches(
                 form the original Series. 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
         """
-        df = self._block.to_pandas_batches(
+        batches = self._block.to_pandas_batches(
             page_size=page_size,
             max_results=max_results,
             allow_large_results=allow_large_results,
-            squeeze=True,
         )
-        return df
+        return map(lambda df: cast(pandas.Series, df.squeeze(1)), batches)
 
     def _compute_dry_run(self) -> bigquery.QueryJob:
         _, query_job = self._block._compute_dry_run((self._value_column,))

tests/system/small/test_dataframe.py

@@ -871,6 +871,21 @@ def test_filter_df(scalars_dfs):
     assert_pandas_df_equal(bf_result, pd_result)
 
 
+def test_df_to_pandas_batches(scalars_dfs):
+    scalars_df, scalars_pandas_df = scalars_dfs
+
+    capped_unfiltered_batches = scalars_df.to_pandas_batches(page_size=2, max_results=6)
+    bf_bool_series = scalars_df["bool_col"]
+    filtered_batches = scalars_df[bf_bool_series].to_pandas_batches()
+
+    pd_bool_series = scalars_pandas_df["bool_col"]
+    pd_result = scalars_pandas_df[pd_bool_series]
+
+    assert 6 == capped_unfiltered_batches.total_rows
+    assert len(pd_result) == filtered_batches.total_rows
+    assert_pandas_df_equal(pd.concat(filtered_batches), pd_result)
+
+
 def test_assign_new_column(scalars_dfs):
     scalars_df, scalars_pandas_df = scalars_dfs
     kwargs = {"new_col": 2}

with_defer.csv

@@ -0,0 +1,45 @@
+Benchmark_Name,Query_Count,Bytes_Processed,Slot_Millis,Local_Execution_Time_Sec,BigQuery_Execution_Time_Sec,Query_Char_Count
+tpch/q1.py_1g_ordered,1.0,276055890.0,,1.11722097202437,,3160.0
+tpch/q1.py_1g_unordered,1.0,276055890.0,,1.1833874560543336,,2693.0
+tpch/q10.py_1g_ordered,1.0,222417387.0,,3.7386660729534924,,6712.0
+tpch/q10.py_1g_unordered,1.0,222417387.0,,3.264317864959594,,6196.0
+tpch/q11.py_1g_ordered,1.0,25760427.0,,2.847489545994904,,7079.0
+tpch/q11.py_1g_unordered,1.0,25760427.0,,2.7726170999812894,,6318.0
+tpch/q12.py_1g_ordered,1.0,257358173.0,,2.054829464002978,,3133.0
+tpch/q12.py_1g_unordered,1.0,257358173.0,,1.8945795510080643,,2884.0
+tpch/q13.py_1g_ordered,1.0,100970808.0,,2.0386499610031024,,1857.0
+tpch/q13.py_1g_unordered,1.0,100970808.0,,1.688826983969193,,1495.0
+tpch/q14.py_1g_ordered,1.0,198158826.0,,2.9061087630107068,,8491.0
+tpch/q14.py_1g_unordered,1.0,198158826.0,,2.597000441979617,,7752.0
+tpch/q15.py_1g_ordered,1.0,192758651.0,,3.323294806003105,,7497.0
+tpch/q15.py_1g_unordered,1.0,192758651.0,,2.1054778510006145,,7100.0
+tpch/q16.py_1g_ordered,1.0,23245641.0,,4.862595856946427,,3883.0
+tpch/q16.py_1g_unordered,1.0,23245641.0,,4.7810218569939025,,3486.0
+tpch/q17.py_1g_ordered,1.0,149544131.0,,2.00097143399762,,5299.0
+tpch/q17.py_1g_unordered,1.0,149544131.0,,1.7620525600505061,,5299.0
+tpch/q18.py_1g_ordered,1.0,148219440.0,,6.1927603500080295,,4092.0
+tpch/q18.py_1g_unordered,1.0,148219440.0,,6.382368930033408,,3645.0
+tpch/q19.py_1g_ordered,1.0,320882154.0,,1.8559044090216048,,4893.0
+tpch/q19.py_1g_unordered,1.0,320882154.0,,1.9689836280304007,,4893.0
+tpch/q2.py_1g_ordered,1.0,31646123.0,,4.538625982997473,,12455.0
+tpch/q2.py_1g_unordered,1.0,31646123.0,,4.537466931971721,,11317.0
+tpch/q20.py_1g_ordered,1.0,220419299.0,,4.036962046986446,,5811.0
+tpch/q20.py_1g_unordered,1.0,220419299.0,,4.330846315948293,,5634.0
+tpch/q21.py_1g_ordered,1.0,208899307.0,,11.05823014699854,,7575.0
+tpch/q21.py_1g_unordered,1.0,208899307.0,,8.225949577987194,,7349.0
+tpch/q22.py_1g_ordered,1.0,16950000.0,,3.3201047479524277,,8010.0
+tpch/q22.py_1g_unordered,1.0,16950000.0,,2.202774106990546,,7320.0
+tpch/q3.py_1g_ordered,1.0,242888490.0,,2.5994084490230307,,3207.0
+tpch/q3.py_1g_unordered,1.0,242888490.0,,4.792361666972283,,2872.0
+tpch/q4.py_1g_ordered,1.0,183628989.0,,1.7768298700102605,,2897.0
+tpch/q4.py_1g_unordered,1.0,183628989.0,,1.8590243599610403,,2676.0
+tpch/q5.py_1g_ordered,1.0,230599591.0,,3.9272409430122934,,5008.0
+tpch/q5.py_1g_unordered,1.0,230599591.0,,4.712351768976077,,4830.0
+tpch/q6.py_1g_ordered,1.0,192038880.0,,1.3108948430162854,,2290.0
+tpch/q6.py_1g_unordered,1.0,192038880.0,,1.2218059509759769,,2203.0
+tpch/q7.py_1g_ordered,1.0,0.0,,3.1901641670265235,,7583.0
+tpch/q7.py_1g_unordered,1.0,0.0,,3.43800652801292,,6757.0
+tpch/q8.py_1g_ordered,1.0,284729257.0,,5.266416695958469,,11675.0
+tpch/q8.py_1g_unordered,1.0,284729257.0,,5.448461340041831,,10892.0
+tpch/q9.py_1g_ordered,1.0,339968968.0,,5.6278779080021195,,8314.0
+tpch/q9.py_1g_unordered,1.0,339968968.0,,5.123836242011748,,7495.0