Skip to content

REF: consolidate ndim checks in _cython_operation #29536

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 1 commit into from
Nov 11, 2019
Merged
Changes from all 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
25 changes: 11 additions & 14 deletions pandas/core/groupby/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,7 @@ def get_group_levels(self):

_cython_arity = {"ohlc": 4} # OHLC

_name_functions = {"ohlc": lambda *args: ["open", "high", "low", "close"]}
_name_functions = {"ohlc": ["open", "high", "low", "close"]}

def _is_builtin_func(self, arg):
"""
Expand Down Expand Up @@ -399,6 +399,13 @@ def _cython_operation(
assert kind in ["transform", "aggregate"]
orig_values = values

if values.ndim > 2:
raise NotImplementedError("number of dimensions is currently limited to 2")
elif values.ndim == 2:
# Note: it is *not* the case that axis is always 0 for 1-dim values,
# as we can have 1D ExtensionArrays that we need to treat as 2D
assert axis == 1, axis

# can we do this operation with our cython functions
# if not raise NotImplementedError

Expand Down Expand Up @@ -524,10 +531,7 @@ def _cython_operation(
if vdim == 1 and arity == 1:
result = result[:, 0]

if how in self._name_functions:
names = self._name_functions[how]() # type: Optional[List[str]]
else:
names = None
names = self._name_functions.get(how, None) # type: Optional[List[str]]

if swapped:
result = result.swapaxes(0, axis)
Expand Down Expand Up @@ -557,10 +561,7 @@ def _aggregate(
is_datetimelike: bool,
min_count: int = -1,
):
if values.ndim > 2:
# punting for now
raise NotImplementedError("number of dimensions is currently limited to 2")
elif agg_func is libgroupby.group_nth:
if agg_func is libgroupby.group_nth:
# different signature from the others
# TODO: should we be using min_count instead of hard-coding it?
agg_func(result, counts, values, comp_ids, rank=1, min_count=-1)
Expand All @@ -574,11 +575,7 @@ def _transform(
):

comp_ids, _, ngroups = self.group_info
if values.ndim > 2:
# punting for now
raise NotImplementedError("number of dimensions is currently limited to 2")
else:
transform_func(result, values, comp_ids, ngroups, is_datetimelike, **kwargs)
transform_func(result, values, comp_ids, ngroups, is_datetimelike, **kwargs)

return result

Expand Down