From a473140473eff4de91c9124b3f52a005db0f1097 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 23 Jan 2024 01:43:14 +0000 Subject: [PATCH 1/2] fix: Series iteration correctly returns values instead of index --- bigframes/series.py | 2 +- tests/system/small/test_series.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/bigframes/series.py b/bigframes/series.py index 2371aad780..e049b41461 100644 --- a/bigframes/series.py +++ b/bigframes/series.py @@ -155,7 +155,7 @@ def __len__(self): def __iter__(self) -> typing.Iterator: return itertools.chain.from_iterable( - map(lambda x: x.index, self._block.to_pandas_batches()) + map(lambda x: x.squeeze(axis=1), self._block.to_pandas_batches()) ) def copy(self) -> Series: diff --git a/tests/system/small/test_series.py b/tests/system/small/test_series.py index 6f919f740f..69b35d102c 100644 --- a/tests/system/small/test_series.py +++ b/tests/system/small/test_series.py @@ -2987,3 +2987,13 @@ def test_sample(scalars_dfs, frac, n, random_state): n = 1 if n is None else n expected_sample_size = round(frac * scalars_df.shape[0]) if frac is not None else n assert bf_result.shape[0] == expected_sample_size + + +def test_series_iter( + scalars_df_index, + scalars_pandas_df_index, +): + for bf_i, pd_i in zip( + scalars_df_index["int64_too"], scalars_pandas_df_index["int64_too"] + ): + assert bf_i == pd_i From f4ae6ae924767277e643fc7cc4ca4686a64d4977 Mon Sep 17 00:00:00 2001 From: Trevor Bergeron Date: Tue, 23 Jan 2024 22:39:52 +0000 Subject: [PATCH 2/2] Update iter docstring --- .../bigframes_vendored/pandas/core/generic.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/third_party/bigframes_vendored/pandas/core/generic.py b/third_party/bigframes_vendored/pandas/core/generic.py index 2ca51f6493..b55c7e23d8 100644 --- a/third_party/bigframes_vendored/pandas/core/generic.py +++ b/third_party/bigframes_vendored/pandas/core/generic.py @@ -50,10 +50,10 @@ def size(self) -> int: def __iter__(self) -> Iterator: """ - Iterate over info axis. + Iterate over column axis for DataFrame, or values for Series. - Returns - iterator: Info axis as iterator. + Returns: + iterator **Examples:** >>> import bigframes.pandas as bpd @@ -71,9 +71,9 @@ def __iter__(self) -> Iterator: >>> series = bpd.Series(["a", "b", "c"], index=[10, 20, 30]) >>> for x in series: ... print(x) - 10 - 20 - 30 + a + b + c """ raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)