feat: add DataFrame.to_pandas_batches() to download large DataFrame objects (#136)

Builds on https://togithub.com/googleapis/python-bigquery-dataframes/pull/132

Towards internal issue 280662868
🦕

Author: tswast

bigframes/core/blocks.py

@@ -416,6 +416,30 @@ def to_pandas(
         )
         return df, query_job
 
+    def to_pandas_batches(self):
+        """Download results one message at a time."""
+        dtypes = dict(zip(self.index_columns, self.index_dtypes))
+        dtypes.update(zip(self.value_columns, self.dtypes))
+        results_iterator, _ = self._expr.start_query()
+        for arrow_table in results_iterator.to_arrow_iterable(
+            bqstorage_client=self._expr._session.bqstoragereadclient
+        ):
+            df = bigframes.session._io.pandas.arrow_to_pandas(arrow_table, dtypes)
+            self._copy_index_to_pandas(df)
+            yield df
+
+    def _copy_index_to_pandas(self, df: pd.DataFrame):
+        """Set the index on pandas DataFrame to match this block.
+
+        Warning: This method modifies ``df`` inplace.
+        """
+        if self.index_columns:
+            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
+
     def _compute_and_count(
         self,
         value_keys: Optional[Iterable[str]] = None,
@@ -489,10 +513,7 @@ def _compute_and_count(
         else:
             total_rows = results_iterator.total_rows
             df = self._to_dataframe(results_iterator)
-
-            if self.index_columns:
-                df.set_index(list(self.index_columns), inplace=True)
-                df.index.names = self.index.names  # type: ignore
+            self._copy_index_to_pandas(df)
 
         return df, total_rows, query_job
 

bigframes/dataframe.py

@@ -893,6 +893,10 @@ def to_pandas(
         self._set_internal_query_job(query_job)
         return df.set_axis(self._block.column_labels, axis=1, copy=False)
 
+    def to_pandas_batches(self) -> Iterable[pandas.DataFrame]:
+        """Stream DataFrame results to an iterable of pandas DataFrame"""
+        return self._block.to_pandas_batches()
+
     def _compute_dry_run(self) -> bigquery.QueryJob:
         return self._block._compute_dry_run()
 

bigframes/session/_io/pandas.py

@@ -46,20 +46,32 @@ def arrow_to_pandas(
             # Preserve NA/NaN distinction. Note: This is currently needed, even if we use
             # nullable Float64Dtype in the types_mapper. See:
             # https://github.com/pandas-dev/pandas/issues/55668
+            mask = pyarrow.compute.is_null(column)
+            nonnull = pyarrow.compute.fill_null(column, float("nan"))
             # Regarding type: ignore, this class has been public at this
             # location since pandas 1.2.0. See:
             # https://pandas.pydata.org/docs/dev/reference/api/pandas.arrays.FloatingArray.html
             pd_array = pandas.arrays.FloatingArray(  # type: ignore
-                column.to_numpy(),
-                pyarrow.compute.is_null(column).to_numpy(),
+                nonnull.to_numpy()
+                if isinstance(nonnull, pyarrow.ChunkedArray)
+                else nonnull.to_numpy(zero_copy_only=False),
+                mask.to_numpy()
+                if isinstance(mask, pyarrow.ChunkedArray)
+                else mask.to_numpy(zero_copy_only=False),
             )
             series = pandas.Series(pd_array, dtype=dtype)
         elif dtype == pandas.Int64Dtype():
             # Avoid out-of-bounds errors in Pandas 1.5.x, which incorrectly
             # casts to float64 in an intermediate step.
+            mask = pyarrow.compute.is_null(column)
+            nonnull = pyarrow.compute.fill_null(column, 0)
             pd_array = pandas.arrays.IntegerArray(
-                pyarrow.compute.fill_null(column, 0).to_numpy(),
-                pyarrow.compute.is_null(column).to_numpy(),
+                nonnull.to_numpy()
+                if isinstance(nonnull, pyarrow.ChunkedArray)
+                else nonnull.to_numpy(zero_copy_only=False),
+                mask.to_numpy()
+                if isinstance(mask, pyarrow.ChunkedArray)
+                else mask.to_numpy(zero_copy_only=False),
             )
             series = pandas.Series(pd_array, dtype=dtype)
         elif isinstance(dtype, pandas.ArrowDtype):

tests/system/small/test_dataframe_io.py

@@ -83,6 +83,14 @@ def test_to_pandas_array_struct_correct_result(session):
     )
 
 
+def test_to_pandas_batches_w_correct_dtypes(scalars_df_default_index):
+    """Verify to_pandas_batches() APIs returns the expected dtypes."""
+    expected = scalars_df_default_index.dtypes
+    for df in scalars_df_default_index.to_pandas_batches():
+        actual = df.dtypes
+        pd.testing.assert_series_equal(actual, expected)
+
+
 @pytest.mark.parametrize(
     ("index"),
     [True, False],

tests/unit/session/test_io_pandas.py

@@ -231,6 +231,62 @@
             ),
             id="scalar-dtypes",
         ),
+        pytest.param(
+            pyarrow.Table.from_pydict(
+                {
+                    "bool": pyarrow.chunked_array(
+                        [[True, None], [True, False]],
+                        type=pyarrow.bool_(),
+                    ),
+                    "bytes": pyarrow.chunked_array(
+                        [[b"123", None], [b"abc", b"xyz"]],
+                        type=pyarrow.binary(),
+                    ),
+                    "float": pyarrow.chunked_array(
+                        [[1.0, None], [float("nan"), -1.0]],
+                        type=pyarrow.float64(),
+                    ),
+                    "int": pyarrow.chunked_array(
+                        [[1, None], [-1, 2**63 - 1]],
+                        type=pyarrow.int64(),
+                    ),
+                    "string": pyarrow.chunked_array(
+                        [["123", None], ["abc", "xyz"]],
+                        type=pyarrow.string(),
+                    ),
+                }
+            ),
+            {
+                "bool": "boolean",
+                "bytes": "object",
+                "float": pandas.Float64Dtype(),
+                "int": pandas.Int64Dtype(),
+                "string": "string[pyarrow]",
+            },
+            pandas.DataFrame(
+                {
+                    "bool": pandas.Series([True, None, True, False], dtype="boolean"),
+                    "bytes": [b"123", None, b"abc", b"xyz"],
+                    "float": pandas.Series(
+                        pandas.arrays.FloatingArray(  # type: ignore
+                            numpy.array(
+                                [1.0, float("nan"), float("nan"), -1.0], dtype="float64"
+                            ),
+                            numpy.array([False, True, False, False], dtype="bool"),
+                        ),
+                        dtype=pandas.Float64Dtype(),
+                    ),
+                    "int": pandas.Series(
+                        [1, None, -1, 2**63 - 1],
+                        dtype=pandas.Int64Dtype(),
+                    ),
+                    "string": pandas.Series(
+                        ["123", None, "abc", "xyz"], dtype="string[pyarrow]"
+                    ),
+                }
+            ),
+            id="scalar-dtypes-chunked_array",
+        ),
         pytest.param(
             pyarrow.Table.from_pydict(
                 {