Skip to content

PERF (CoW): optimize equality check in groupby is_in_obj #50730

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.

Sign up for GitHub

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

Closed
wants to merge 1 commit into from
Closed

PERF (CoW): optimize equality check in groupby is_in_obj #50730

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
31 changes: 30 additions & 1 deletion pandas/core/groupby/grouper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,18 @@

from pandas.core.dtypes.common import (
is_categorical_dtype,
is_dtype_equal,
is_list_like,
is_scalar,
)

from pandas.core import algorithms
from pandas.core.arrays import (
BaseMaskedArray,
Categorical,
ExtensionArray,
)
from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
import pandas.core.common as com
from pandas.core.frame import DataFrame
from pandas.core.groupby import ops
Expand Down Expand Up @@ -893,7 +896,7 @@ def is_in_obj(gpr) -> bool:
# For the CoW case, we need an equality check as the identity check
# no longer works (each Series from column access is a new object)
try:
return gpr.equals(obj[gpr.name])
return _is_same_memory(gpr, obj[gpr.name])
except (AttributeError, KeyError, IndexError, InvalidIndexError):
return False
try:
Expand Down Expand Up @@ -986,3 +989,29 @@ def _convert_grouper(axis: Index, grouper):
return grouper
else:
return grouper


def _is_same_memory(s1, s2):
if not is_dtype_equal(s1, s2):
return False
if len(s1) != len(s2):
return False

arr1 = s1._values
arr2 = s2._values

if isinstance(arr1, np.ndarray):
# slice first element -> if first element shares memory and they are
# equal length -> share exactly memory for the full length (if not
# checking first element, two arrays might overlap partially)
Copy link
Member

Choose a reason for hiding this comment

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

are they necessarily contiguous?

Copy link
Member Author

Choose a reason for hiding this comment

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

That's a good question (and we should cover that in tests). No, not necessarily I suppose. But the question is if that matters?

In theory they could have a first same memory point, but if not contiguous and if they have different strides, they might still have different values.
But can you get that in practice with a Series from a DataFrame?

Copy link
Member Author

Choose a reason for hiding this comment

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

Can of course easily check that the strides are equal. But is that enough guarantee?

Copy link
Member

@phofl phofl Jan 13, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

You can get there, probably also in the groupby case:

df = DataFrame({"a": list(range(100)), "b": 1})

ser = df.loc[slice(0, 50), "a"]
ser2 = df.loc[slice(0, 100, 2), "a"]
np.shares_memory(ser._values, ser2._values)

Edit: checking the strides fixes this of course

return np.shares_memory(arr1[:1], arr2[:1])

elif isinstance(arr1, BaseMaskedArray):
return np.shares_memory(arr1._data[:1], arr2._data[:1]) and np.shares_memory(
arr1._mask[:1], arr2._mask[:1]
)

elif isinstance(arr1, NDArrayBackedExtensionArray):
return np.shares_memory(arr1._ndarray[:1], arr2._ndarray[:1])
else:
return s1.equals(s2)