Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
44 commits
Select commit Hold shift + click to select a range
7cd4dc3
wip
mzeitlin11 Mar 26, 2021
91984dc
wip
mzeitlin11 Mar 26, 2021
b371cc5
wip
mzeitlin11 Mar 26, 2021
69cce96
wip
mzeitlin11 Mar 26, 2021
dd7f324
wip
mzeitlin11 Mar 26, 2021
64680d4
wip
mzeitlin11 Mar 26, 2021
be16f65
wip
mzeitlin11 Mar 26, 2021
9442846
wip
mzeitlin11 Mar 26, 2021
f089175
wip
mzeitlin11 Mar 26, 2021
31409f8
wip
mzeitlin11 Mar 26, 2021
0c05f74
wip
mzeitlin11 Mar 26, 2021
18dcc94
Merge remote-tracking branch 'origin/master' into perf/masked_cummin/max
mzeitlin11 Mar 26, 2021
5c60a1f
wip
mzeitlin11 Mar 26, 2021
f0c27ce
PERF: use masked algo in groupby cummin and cummax
mzeitlin11 Mar 27, 2021
2fa80ad
Avoid mask copy
mzeitlin11 Mar 27, 2021
280c7e5
Update whatsnew
mzeitlin11 Mar 27, 2021
dca28cf
Merge remote-tracking branch 'origin/master' into perf/masked_cummin/max
mzeitlin11 Apr 1, 2021
7e2fbe0
Merge fixup
mzeitlin11 Apr 1, 2021
0ebb97a
Follow transpose
mzeitlin11 Apr 1, 2021
0009dfd
Compute mask usage inside algo
mzeitlin11 Apr 1, 2021
6663832
try optional
mzeitlin11 Apr 1, 2021
8247f82
WIP
mzeitlin11 Apr 1, 2021
71e1c4f
Use more contiguity
mzeitlin11 Apr 1, 2021
c6cf9ee
Shrink benchmark
mzeitlin11 Apr 1, 2021
02768ec
Merge remote-tracking branch 'origin/master' into perf/masked_cummin/max
mzeitlin11 Apr 1, 2021
836175b
Merge remote-tracking branch 'origin/master' into perf/masked_cummin/max
mzeitlin11 Apr 2, 2021
293dc6e
Revert unrelated
mzeitlin11 Apr 2, 2021
478c6c9
Merge remote-tracking branch 'origin/master' into perf/masked_cummin/max
mzeitlin11 Apr 6, 2021
fa45a9a
Merge remote-tracking branch 'origin/master' into perf/masked_cummin/max
mzeitlin11 Apr 8, 2021
1632b81
Merge remote-tracking branch 'origin/master' into perf/masked_cummin/max
mzeitlin11 Apr 12, 2021
1bb344e
Remove merge conflict relic
mzeitlin11 Apr 12, 2021
97d9eea
Update doc/source/whatsnew/v1.3.0.rst
mzeitlin11 Apr 13, 2021
892a92a
Update doc/source/whatsnew/v1.3.0.rst
mzeitlin11 Apr 13, 2021
a239a68
Update pandas/core/groupby/ops.py
mzeitlin11 Apr 13, 2021
f98ca35
Merge remote-tracking branch 'origin' into perf/masked_cummin/max
mzeitlin11 Apr 13, 2021
e7ed12f
Merge branch 'perf/masked_cummin/max' of github.com:/mzeitlin11/panda…
mzeitlin11 Apr 13, 2021
a1422ba
Address comments
mzeitlin11 Apr 13, 2021
482a209
Change random generation style
mzeitlin11 Apr 13, 2021
4e7404d
Merge remote-tracking branch 'origin' into perf/masked_cummin/max
mzeitlin11 Apr 18, 2021
251c02a
Use conditional instead of partial
mzeitlin11 Apr 18, 2021
3de7e5e
Remove ensure_int_or_float
mzeitlin11 Apr 18, 2021
237f86f
Remove unnecessary condition
mzeitlin11 Apr 18, 2021
a1b0c04
Merge remote-tracking branch 'origin' into perf/masked_cummin/max
mzeitlin11 Apr 19, 2021
5e1dac4
Merge remote-tracking branch 'origin' into perf/masked_cummin/max
mzeitlin11 Apr 20, 2021
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
Prev Previous commit
Next Next commit
wip
  • Loading branch information
mzeitlin11 committed Mar 26, 2021
commit 5c60a1f35888920853a324353c9ccad90cf62e11
2 changes: 0 additions & 2 deletions asv_bench/benchmarks/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,10 +395,8 @@ class GroupByMethods:
params = [
["int", "float", "object", "datetime"],
[

"cummax",
"cummin",

],
["direct", "transformation"],
]
Expand Down
32 changes: 24 additions & 8 deletions pandas/_libs/groupby.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -1270,14 +1270,10 @@ def group_cummin_max(groupby_t[:, ::1] out,
Array to store cummin/max in.
values : array
Values to take cummin/max of.
<<<<<<< HEAD
mask : array[uint8_t]
If `use_mask`, then indices represent missing values,
otherwise will be passed as a zeroed array
labels : int64 array
=======
labels : np.ndarray[np.intp]
>>>>>>> origin/master
Labels to group by.
ngroups : int
Number of groups, larger than all entries of `labels`.
Expand Down Expand Up @@ -1320,14 +1316,16 @@ def group_cummin_max(groupby_t[:, ::1] out,
for j in range(K):
val_is_nan = False

# If using the mask, we can avoid grabbing the
# value unless necessary
if use_mask:
if mask[i, j]:

# `out` does not need to be set since it
# will be masked anyway
val_is_nan = True
else:

# If using the mask, we can avoid grabbing the
# value unless necessary
val = values[i, j]

# Otherwise, `out` must be set accordingly if the
Expand Down Expand Up @@ -1359,7 +1357,16 @@ def group_cummin(groupby_t[:, ::1] out,
bint is_datetimelike,
bint use_mask):
"""See group_cummin_max.__doc__"""
group_cummin_max(out, values, mask, labels, ngroups, is_datetimelike, use_mask, compute_max=False)
group_cummin_max(
out,
values,
mask,
labels,
ngroups,
is_datetimelike,
use_mask,
compute_max=False
)


@cython.boundscheck(False)
Expand All @@ -1372,4 +1379,13 @@ def group_cummax(groupby_t[:, ::1] out,
bint is_datetimelike,
bint use_mask):
"""See group_cummin_max.__doc__"""
group_cummin_max(out, values, mask, labels, ngroups, is_datetimelike, use_mask, compute_max=True)
group_cummin_max(
out,
values,
mask,
labels,
ngroups,
is_datetimelike,
use_mask,
compute_max=True
)
2 changes: 0 additions & 2 deletions pandas/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,6 @@ class providing the base-class of operations.
from pandas.core.arrays import (
Categorical,
ExtensionArray,
BaseMaskedArray
)
from pandas.core.base import (
DataError,
Expand All @@ -105,7 +104,6 @@ class providing the base-class of operations.
Index,
MultiIndex,
)
from pandas.core.groupby.ops import does_cython_function_use_mask
from pandas.core.series import Series
from pandas.core.sorting import get_group_index_sorter
from pandas.core.util.numba_ import NUMBA_FUNC_CACHE
Expand Down
83 changes: 59 additions & 24 deletions pandas/core/groupby/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,6 @@
Tuple,
Type,
)
from pandas.core.arrays.masked import (
BaseMaskedDtype,
)

import numpy as np

Expand All @@ -40,7 +37,6 @@
FrameOrSeries,
Shape,
final,
ArrayLike
)
from pandas.errors import AbstractMethodError
from pandas.util._decorators import cache_readonly
Expand Down Expand Up @@ -75,10 +71,11 @@
isna,
maybe_fill,
)
from pandas.core.arrays import (
BaseMaskedArray
)

from pandas.core.arrays.masked import (
BaseMaskedArray,
BaseMaskedDtype,
)
from pandas.core.base import SelectionMixin
import pandas.core.common as com
from pandas.core.frame import DataFrame
Expand Down Expand Up @@ -123,7 +120,11 @@
"cummax": "group_cummax",
"rank": "group_rank",
},
"needs_mask": {"cummin", "cummax"}
}

_CYTHON_MASKED_FUNCTIONS = {
"cummin",
"cummax",
}


Expand Down Expand Up @@ -163,8 +164,8 @@ def _get_cython_function(kind: str, how: str, dtype: np.dtype, is_numeric: bool)
return func


def does_cython_function_use_mask(kind: str) -> bool:
return kind in _CYTHON_FUNCTIONS["needs_mask"]
def cython_function_uses_mask(kind: str) -> bool:
return kind in _CYTHON_MASKED_FUNCTIONS


class BaseGrouper:
Expand Down Expand Up @@ -590,8 +591,14 @@ def _ea_wrap_cython_operation(

@final
def _masked_ea_wrap_cython_operation(
self, kind: str, values, how: str, axis: int, min_count: int = -1, **kwargs
) -> ArrayLike:
self,
kind: str,
values: BaseMaskedArray,
how: str,
axis: int,
min_count: int = -1,
**kwargs,
) -> BaseMaskedArray:
"""
Equivalent of `_ea_wrap_cython_operation`, but optimized for masked EA's
and cython algorithms which accept a mask.
Expand All @@ -601,23 +608,33 @@ def _masked_ea_wrap_cython_operation(
# isna just directly returns self._mask, so copy here to prevent
# modifying the original
mask = isna(values).copy()
values = values._data
arr = values._data

if is_integer_dtype(values.dtype) or is_bool_dtype(values.dtype):
Copy link
Contributor

Choose a reason for hiding this comment

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

why is this an entirely different funtion? pls integrate to existing infra

Copy link
Member

Choose a reason for hiding this comment

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

This is a different function because it's specific for MaskedArrays. Having this as a separate function is consistent with how it's currently implemented IMO (with a similar separate function for generic EAs).

It could also be another elif check in _ea_wrap_cython_operation, but it's not that it would result in less code or so, and since _ea_wrap_cython_operation already gets quite complicated, I think this separate function is good.

Copy link
Member Author

Choose a reason for hiding this comment

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

In an initial version I tried to fold this into _ea_wrap_cython_operation, but thought this smaller function was a cleaner solution since the conditionals here can remain much simpler. While adding a function for just supporting masked cummax, cummin seems wasteful, this infrastructure should extend to more masked groupby algos.

Copy link
Contributor

Choose a reason for hiding this comment

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

see my comments. This MUST integrate with the existing infrastructure (or refactor that). Duplicating is -1

Copy link
Member

Choose a reason for hiding this comment

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

  1. The whole point of this PR is to add a special handling for masked arrays, to ensure to pass through the mask to the groupby cython algo. This will always add some code (and the code in _masked_ea_wrap_cython_operation is specific to masked arrays).
  2. IMO this is integrated with the existing infrastructure: it integrates nicely into the existing _cython_operation and follows the same pattern as we already have for _ea_wrap_cython_operation

If you don't like how the added code is structured right now, please do a concrete suggestion of how you would do it differently.

# IntegerArray or BooleanArray
values = ensure_int_or_float(values)
arr = ensure_int_or_float(arr)
Copy link
Member

Choose a reason for hiding this comment

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

FWIW im planning to kill off this function; for EAs this is always just arr.to_numpy(dtype="float64", na_value=np.nan)

Copy link
Member Author

Choose a reason for hiding this comment

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

Thanks for bringing up - realized this whole condition can be simplified since we actually have an ndarray at this point


res_values = self._cython_operation(
kind, values, how, axis, min_count, mask=mask, **kwargs
kind, arr, how, axis, min_count, mask=mask, **kwargs
)
dtype = maybe_cast_result_dtype(orig_values.dtype, how)
assert isinstance(dtype, BaseMaskedDtype)
cls = dtype.construct_array_type()

return cls(res_values.astype(dtype.type, copy=False), mask.astype(bool, copy=True))
return cls(
res_values.astype(dtype.type, copy=False), mask.astype(bool, copy=True)
)

@final
def _cython_operation(
self, kind: str, values, how: str, axis: int, min_count: int = -1, mask: Optional[np.ndarray] = None, **kwargs
self,
kind: str,
values,
how: str,
axis: int,
min_count: int = -1,
mask: np.ndarray | None = None,
**kwargs,
) -> ArrayLike:
"""
Returns the values of a cython operation.
Expand All @@ -640,7 +657,7 @@ def _cython_operation(
self._disallow_invalid_ops(dtype, how, is_numeric)

if is_extension_array_dtype(dtype):
if isinstance(dtype, BaseMaskedDtype) and does_cython_function_use_mask(how):
if isinstance(values, BaseMaskedArray) and cython_function_uses_mask(how):
return self._masked_ea_wrap_cython_operation(
kind, values, how, axis, min_count, **kwargs
Copy link
Contributor

Choose a reason for hiding this comment

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

i really don't understand all of this code duplication. this is adding huge complexity. pls reduce it.

Copy link
Member

@jorisvandenbossche jorisvandenbossche Apr 13, 2021

Choose a reason for hiding this comment

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

Jeff, did you actually read the previous responses to your similar comment? (https://github.com/pandas-dev/pandas/pull/40651/files#r603319910) Can you then please answer there to the concrete reasons given.

Copy link
Contributor

Choose a reason for hiding this comment

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

yes and its a terrible pattern.

Copy link
Contributor

Choose a reason for hiding this comment

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

this duplication of code is ridiculous. We have a VERY large codebase. Having this kind of separate logic is amazingling confusing & is humungous tech debt. This is heavily used code and needs to be carefully modified.

Copy link
Member Author

@mzeitlin11 mzeitlin11 Apr 13, 2021

Choose a reason for hiding this comment

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

I understand the concern about adding code complexity - my thinking was that if the goal is for nullable types to become the default in pandas, then direct support makes sense. And in that case, nullable types would need to be special-cased somewhere, and I think the separate function is cleaner than interleaving in _ea_wrap_cython_operation.

If direct support for nullable dtypes is not desired, we can just close this. If it is, I'll keep trying to think of ways to achieve this without adding more code, but any suggestions there would be welcome!

Copy link
Member

@jorisvandenbossche jorisvandenbossche Apr 13, 2021

Choose a reason for hiding this comment

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

Proper support for nullable dtypes is certainly desired (how to add it exactly can of course be discussed), so thanks a lot @mzeitlin11 for your efforts here.

AFAIK, it's correct we need some special casing for it somewhere (that's the whole point of this PR is to add special handling for it).
Where exactly to put this special casing can of course be discussed, but to me the separate helper method instead of interleaving it in _ea_wrap_cython_operation seems good (I don't think that interleaving it into the existing _ea_wrap_cython_operation would result in fewer added lines of code (and would be harder to read)).

@jreback please try to stay constructive (eg answer to our arguments or provide concrete suggestions on where you would put it / how you would do it differently) and please mind your language (there is no need to call the approach taken by a contributor "terrible").

Copy link
Member

Choose a reason for hiding this comment

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

  1. I agree with @jorisvandenbossche on phrasing concerns. Even the best of us slip up here from time to time.

  2. if the goal is for nullable types to become the default in pandas

This decision has not been made.

  1. I think the separate function is cleaner than interleaving in _ea_wrap_cython_operation.

Agreed.

  1. My preferred dispatch logic would look something like:
def _cython_operation(...)
    if is_ea_dtype(...):
       return self. _ea_wrap_cython_operation(...)
    [status quo]

def _ea_wrap_cython_operation(...):
    if should_use_mask(...):
        return self._masked_ea_wrap_cython_operation(...)
    [status quo]

as Joris correctly pointed out, that is not viable ATM. I think a lot of this dispatch logic eventually belongs in WrappedCythonOp (which i've been vaguely planning on doing next time there aren't any open PRs touching this code), at which point we can reconsider flattening this

  1. My other preferred dispatch logic would not be in this file at all, but be implemented as a method on the EA subclass. I'm really uncomfortable with this code depending on MaskedArray implementation details, seeing as how there has been discussion of swapping them out for something arrow-based.

Copy link
Member Author

Choose a reason for hiding this comment

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

@jbrockmendel if you plan further refactoring of this code, I'm happy to just mothball this pr for now. The real benefit won't come in until more groupby algos allow a mask on this path anyway, so not worth adding now if it's just going to cause more pain in future refactoring.

I also like the idea of approach 5 instead of this - could start looking into that if you think it's a promising direction.

Copy link
Member

Choose a reason for hiding this comment

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

if you plan further refactoring of this code, I'm happy to just mothball this pr for now.

From today's call, I think the plan is to move forward with this first.

I also like the idea of approach 5 instead of this - could start looking into that if you think it's a promising direction.

Long-term I think this is the right way to go to get the general case right, so I'd encourage you if you're interested in trying to implement this on the EA- separate PR(s).

)
Expand Down Expand Up @@ -689,7 +706,9 @@ def _cython_operation(
)
out_shape = (self.ngroups,) + values.shape[1:]

func, values, needs_mask = self._get_cython_func_and_vals(kind, how, values, is_numeric)
func, values, needs_mask = self._get_cython_func_and_vals(
kind, how, values, is_numeric
)
use_mask = mask is not None
if needs_mask:
if mask is None:
Expand All @@ -716,10 +735,10 @@ def _cython_operation(
)

if not use_mask and is_integer_dtype(result.dtype) and not is_datetimelike:
mask = result == iNaT
if mask.any():
result_mask = result == iNaT
if result_mask.any():
result = result.astype("float64")
result[mask] = np.nan
result[result_mask] = np.nan

if kind == "aggregate" and self._filter_empty_groups and not counts.all():
assert result.ndim != 2
Expand Down Expand Up @@ -755,12 +774,28 @@ def _aggregate(

@final
def _transform(
self, result: np.ndarray, values: np.ndarray, transform_func, is_datetimelike: bool, use_mask: bool, mask: np.ndarray | None, **kwargs
self,
result: np.ndarray,
values: np.ndarray,
transform_func,
is_datetimelike: bool,
use_mask: bool,
mask: np.ndarray | None,
**kwargs,
) -> np.ndarray:

comp_ids, _, ngroups = self.group_info
if mask is not None:
transform_func(result, values, mask, comp_ids, ngroups, is_datetimelike, use_mask, **kwargs)
transform_func(
result,
values,
mask,
comp_ids,
ngroups,
is_datetimelike,
use_mask,
**kwargs,
)
else:
transform_func(result, values, comp_ids, ngroups, is_datetimelike, **kwargs)

Expand Down
12 changes: 8 additions & 4 deletions pandas/tests/groupby/test_function.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,14 @@ def dtypes_for_minmax(request):
np_type = np.float64

min_val = (
np.iinfo(np_type).min if np.dtype(np_type).kind == "i" else np.finfo(np_type).min
np.iinfo(np_type).min
if np.dtype(np_type).kind == "i"
else np.finfo(np_type).min
)
max_val = (
np.iinfo(np_type).max if np.dtype(np_type).kind == "i" else np.finfo(np_type).max
np.iinfo(np_type).max
if np.dtype(np_type).kind == "i"
else np.finfo(np_type).max
)

return (dtype, min_val, max_val)
Expand Down Expand Up @@ -855,11 +859,11 @@ def test_cummax(dtypes_for_minmax):
)
def test_nullable_int_not_cast_as_float(method, dtype, val):
data = [val, pd.NA]
df = pd.DataFrame({"grp": [1, 1], "b": data}, dtype=dtype)
df = DataFrame({"grp": [1, 1], "b": data}, dtype=dtype)
grouped = df.groupby("grp")

result = grouped.transform(method)
expected = pd.DataFrame({"b": data}, dtype=dtype)
expected = DataFrame({"b": data}, dtype=dtype)

tm.assert_frame_equal(result, expected)

Expand Down