Skip to content

REF: reverse dispatch in factorize #49966

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
Dec 1, 2022
Merged
Show file tree
Hide file tree
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
24 changes: 3 additions & 21 deletions pandas/core/algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,6 @@
ABCExtensionArray,
ABCIndex,
ABCMultiIndex,
ABCRangeIndex,
ABCSeries,
ABCTimedeltaArray,
)
Expand Down Expand Up @@ -738,13 +737,11 @@ def factorize(
# Step 2 is dispatched to extension types (like Categorical). They are
# responsible only for factorization. All data coercion, sorting and boxing
# should happen here.
if isinstance(values, ABCRangeIndex):
return values.factorize(sort=sort)
if isinstance(values, (ABCIndex, ABCSeries)):
return values.factorize(sort=sort, use_na_sentinel=use_na_sentinel)

values = _ensure_arraylike(values)
original = values
if not isinstance(values, ABCMultiIndex):
values = extract_array(values, extract_numpy=True)
Comment on lines -746 to -747
Copy link
Member

Choose a reason for hiding this comment

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

is this no longer necessary because of the early return above

   if isinstance(values, (ABCIndex, ABCSeries)):
        return values.factorize(sort=sort, use_na_sentinel=use_na_sentinel)

?

Copy link
Member Author

Choose a reason for hiding this comment

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

correct


if (
isinstance(values, (ABCDatetimeArray, ABCTimedeltaArray))
Expand All @@ -753,7 +750,7 @@ def factorize(
# The presence of 'freq' means we can fast-path sorting and know there
# aren't NAs
codes, uniques = values.factorize(sort=sort)
return _re_wrap_factorize(original, uniques, codes)
return codes, uniques

elif not isinstance(values.dtype, np.dtype):
codes, uniques = values.factorize(use_na_sentinel=use_na_sentinel)
Expand Down Expand Up @@ -789,21 +786,6 @@ def factorize(

uniques = _reconstruct_data(uniques, original.dtype, original)

return _re_wrap_factorize(original, uniques, codes)


def _re_wrap_factorize(original, uniques, codes: np.ndarray):
"""
Wrap factorize results in Series or Index depending on original type.
"""
if isinstance(original, ABCIndex):
uniques = ensure_wrapped_if_datetimelike(uniques)
uniques = original._shallow_copy(uniques, name=None)
elif isinstance(original, ABCSeries):
from pandas import Index

uniques = Index(uniques)

return codes, uniques


Expand Down
17 changes: 15 additions & 2 deletions pandas/core/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@

from pandas import (
Categorical,
Index,
Series,
)

Expand Down Expand Up @@ -1134,8 +1135,20 @@ def factorize(
self,
sort: bool = False,
use_na_sentinel: bool = True,
):
return algorithms.factorize(self, sort=sort, use_na_sentinel=use_na_sentinel)
) -> tuple[npt.NDArray[np.intp], Index]:

codes, uniques = algorithms.factorize(
self._values, sort=sort, use_na_sentinel=use_na_sentinel
)

if isinstance(self, ABCIndex):
# preserve e.g. NumericIndex, preserve MultiIndex
uniques = self._constructor(uniques)
Copy link
Member

Choose a reason for hiding this comment

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

what's the reason to change this from

            uniques = ensure_wrapped_if_datetimelike(uniques)
            uniques = self._shallow_copy(uniques, name=None)

? Is it just simpler?

Copy link
Member Author

Choose a reason for hiding this comment

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

correct

else:
Copy link
Member

Choose a reason for hiding this comment

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

does this change from elif isinstance(self, ABCSeries) because here self would always be ABCSeries, so it might as well just be else?

Copy link
Member Author

Choose a reason for hiding this comment

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

correct

from pandas import Index

uniques = Index(uniques)
return codes, uniques

_shared_docs[
"searchsorted"
Expand Down