Skip to content

Commit

Permalink
feat: support series items method (#1089)
Browse files Browse the repository at this point in the history
* feat: series items method

* fix lint format

* use to_pandas_batches

* fix doctest
  • Loading branch information
jialuoo authored Oct 23, 2024
1 parent e8d689e commit 245a89c
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 0 deletions.
8 changes: 8 additions & 0 deletions bigframes/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -1117,6 +1117,14 @@ def ne(self, other: object) -> Series:
# TODO: enforce stricter alignment
return self._apply_binary_op(other, ops.ne_op)

def items(self):
for batch_df in self._block.to_pandas_batches():
assert (
batch_df.shape[1] == 1
), f"Expected 1 column in the dataframe, but got {batch_df.shape[1]}."
for item in batch_df.squeeze(axis=1).items():
yield item

def where(self, cond, other=None):
value_id, cond_id, other_id, block = self._align3(cond, other)
block, result_id = block.project_expr(
Expand Down
24 changes: 24 additions & 0 deletions tests/system/small/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,30 @@ def test_series_construct_geodata():
)


@pytest.mark.parametrize(
["data", "index"],
[
(["a", "b", "c"], None),
([1, 2, 3], ["a", "b", "c"]),
([1, 2, None], ["a", "b", "c"]),
([1, 2, 3], [pd.NA, "b", "c"]),
([numpy.nan, 2, 3], ["a", "b", "c"]),
],
)
def test_series_items(data, index):
bf_series = series.Series(data, index=index)
pd_series = pd.Series(data, index=index)

for (bf_index, bf_value), (pd_index, pd_value) in zip(
bf_series.items(), pd_series.items()
):
# TODO(jialuo): Remove the if conditions after b/373699458 is addressed.
if not pd.isna(bf_index) or not pd.isna(pd_index):
assert bf_index == pd_index
if not pd.isna(bf_value) or not pd.isna(pd_value):
assert bf_value == pd_value


@pytest.mark.parametrize(
["col_name", "expected_dtype"],
[
Expand Down
36 changes: 36 additions & 0 deletions third_party/bigframes_vendored/pandas/core/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -3332,6 +3332,42 @@ def kurt(self):
"""
raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)

def items(self):
"""
Iterate over (index, value) pairs of a Series.
Iterates over the Series contents, returning a tuple with
the index and the value of a Series.
**Examples:**
>>> import bigframes.pandas as bpd
>>> bpd.options.display.progress_bar = None
>>> s = bpd.Series(['bear', 'bear', 'marsupial'],
... index=['panda', 'polar', 'koala'])
>>> s
panda bear
polar bear
koala marsupial
dtype: string
>>> for index, value in s.items():
... print(f'--> index: {index}')
... print(f'--> value: {value}')
...
--> index: panda
--> value: bear
--> index: polar
--> value: bear
--> index: koala
--> value: marsupial
Returns:
Iterator: Iterator of index, value for each content of the Series.
"""
raise NotImplementedError(constants.ABSTRACT_METHOD_ERROR_MESSAGE)

def where(self, cond, other):
"""Replace values where the condition is False.
Expand Down

0 comments on commit 245a89c

Please sign in to comment.