Skip to content

feat: define list accessor for bigframes Series #946

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Sep 5, 2024
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions bigframes/operations/_op_converters.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import bigframes.operations as ops


def convert_index(key: int) -> ops.ArrayIndexOp:
if key < 0:
raise NotImplementedError("Negative indexing is not supported.")
return ops.ArrayIndexOp(index=key)


def convert_slice(key: slice) -> ops.ArraySliceOp:
if key.step is not None and key.step != 1:
raise NotImplementedError(f"Only a step of 1 is allowed, got {key.step}")

if (key.start is not None and key.start < 0) or (
key.stop is not None and key.stop < 0
):
raise NotImplementedError("Slicing with negative numbers is not allowed.")

return ops.ArraySliceOp(
start=key.start if key.start is not None else 0,
stop=key.stop,
step=key.step,
)
46 changes: 46 additions & 0 deletions bigframes/operations/lists.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import inspect
from typing import Union

import bigframes_vendored.pandas.core.arrays.arrow.accessors as vendoracessors

from bigframes.core import log_adapter
import bigframes.operations as ops
from bigframes.operations._op_converters import convert_index, convert_slice
import bigframes.operations.base
import bigframes.series as series


@log_adapter.class_logger
class ListAccessor(
bigframes.operations.base.SeriesMethods, vendoracessors.ListAccessor
):
__doc__ = vendoracessors.ListAccessor.__doc__

def len(self):
return self._apply_unary_op(ops.len_op)

def __getitem__(self, key: Union[int, slice]) -> series.Series:
if isinstance(key, int):
return self._apply_unary_op(convert_index(key))
elif isinstance(key, slice):
return self._apply_unary_op(convert_slice(key))
else:
raise ValueError(f"key must be an int or slice, got {type(key).__name__}")

__getitem__.__doc__ = inspect.getdoc(vendoracessors.ListAccessor.__getitem__)
24 changes: 3 additions & 21 deletions bigframes/operations/strings.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from bigframes.core import log_adapter
import bigframes.dataframe as df
import bigframes.operations as ops
from bigframes.operations._op_converters import convert_index, convert_slice
import bigframes.operations.base
import bigframes.series as series

Expand All @@ -40,28 +41,9 @@ class StringMethods(bigframes.operations.base.SeriesMethods, vendorstr.StringMet

def __getitem__(self, key: Union[int, slice]) -> series.Series:
if isinstance(key, int):
if key < 0:
raise NotImplementedError("Negative indexing is not supported.")
return self._apply_unary_op(ops.ArrayIndexOp(index=key))
return self._apply_unary_op(convert_index(key))
elif isinstance(key, slice):
if key.step is not None and key.step != 1:
raise NotImplementedError(
f"Only a step of 1 is allowed, got {key.step}"
)
if (key.start is not None and key.start < 0) or (
key.stop is not None and key.stop < 0
):
raise NotImplementedError(
"Slicing with negative numbers is not allowed."
)

return self._apply_unary_op(
ops.ArraySliceOp(
start=key.start if key.start is not None else 0,
stop=key.stop,
step=key.step,
)
)
return self._apply_unary_op(convert_slice(key))
else:
raise ValueError(f"key must be an int or slice, got {type(key).__name__}")

Expand Down
9 changes: 8 additions & 1 deletion bigframes/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
import bigframes.operations.aggregations as agg_ops
import bigframes.operations.base
import bigframes.operations.datetimes as dt
import bigframes.operations.lists as lists
import bigframes.operations.plotting as plotting
import bigframes.operations.strings as strings
import bigframes.operations.structs as structs
Expand All @@ -66,6 +67,8 @@
" Try converting it to a remote function."
)

_list = list # Type alias to escape Series.list property


@log_adapter.class_logger
class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Series):
Expand Down Expand Up @@ -161,6 +164,10 @@ def query_job(self) -> Optional[bigquery.QueryJob]:
def struct(self) -> structs.StructAccessor:
return structs.StructAccessor(self._block)

@property
def list(self) -> lists.ListAccessor:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, please add the new accessor to docs/reference/bigframes.pandas/series.rst

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Change updated. PTAL!

return lists.ListAccessor(self._block)

@property
@validations.requires_ordering()
def T(self) -> Series:
Expand Down Expand Up @@ -1708,7 +1715,7 @@ def to_latex(
buf, columns=columns, header=header, index=index, **kwargs
)

def tolist(self) -> list:
def tolist(self) -> _list:
return self.to_pandas().to_list()

to_list = tolist
Expand Down
8 changes: 8 additions & 0 deletions docs/reference/bigframes.pandas/series.rst
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,14 @@ String handling
:inherited-members:
:undoc-members:

List handling
^^^^^^^^^^^^^

.. automodule:: bigframes.operations.lists
:members:
:inherited-members:
:undoc-members:

Struct handling
^^^^^^^^^^^^^^^

Expand Down
Loading