Skip to content

BUG: preserve freq in DTI/TDI factorize #33836

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 10 commits into from
Closed

BUG: preserve freq in DTI/TDI factorize #33836

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
6d741b1
BUG: preserve freq in DTI/TDI factorize
jbrockmendel Apr 28, 2020
e466d71
Merge branch 'master' of https://github.com/pandas-dev/pandas into fr…
jbrockmendel Apr 28, 2020
a553174
mypy fixup
jbrockmendel Apr 28, 2020
23911ef
dummy commit to force CI
jbrockmendel Apr 28, 2020
516d232
Merge branch 'master' of https://github.com/pandas-dev/pandas into fr…
jbrockmendel Apr 28, 2020
0e51930
refactor per joris suggestion
jbrockmendel Apr 28, 2020
678251d
32bit compat
jbrockmendel Apr 28, 2020
abb5913
return copy
jbrockmendel Apr 28, 2020
c96c1ac
Merge branch 'master' of https://github.com/pandas-dev/pandas into fr…
jbrockmendel Apr 29, 2020
7c66389
preserve freq in pd.factorize
jbrockmendel Apr 29, 2020
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
13 changes: 13 additions & 0 deletions pandas/core/algorithms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,11 +46,15 @@
needs_i8_conversion,
)
from pandas.core.dtypes.generic import (
ABCDatetimeArray,
ABCDatetimeIndex,
ABCExtensionArray,
ABCIndex,
ABCIndexClass,
ABCMultiIndex,
ABCSeries,
ABCTimedeltaArray,
ABCTimedeltaIndex,
)
from pandas.core.dtypes.missing import isna, na_value_for_dtype

Expand Down Expand Up @@ -614,6 +618,15 @@ def factorize(
values = _ensure_arraylike(values)
original = values

if isinstance(
values,
(ABCDatetimeIndex, ABCTimedeltaIndex, ABCDatetimeArray, ABCTimedeltaArray),
):
# Defer to the method in order to retain freq
if sort is False and size_hint is None:
# EA/Index methods dont support the same kwargs as this func
return values.factorize(na_sentinel=na_sentinel)

if is_extension_array_dtype(values.dtype):
values = extract_array(values)
codes, uniques = values.factorize(na_sentinel=na_sentinel)
Expand Down
8 changes: 8 additions & 0 deletions pandas/core/arrays/datetimelike.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -437,6 +437,14 @@ def _with_freq(self, freq):
arr._freq = freq
return arr

def factorize(self, na_sentinel=-1):
Copy link
Contributor

Choose a reason for hiding this comment

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

we don't want to expose factorize on Index at all.

Copy link
Member Author

Choose a reason for hiding this comment

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

It is already there in the status quo. If you want to remove it, that needs a deprecation cycle, is a separate issue

Copy link
Contributor

Choose a reason for hiding this comment

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

/Users/jreback/pandas
bash-3.2$ grep -r factorize pandas/core/indexes/
pandas/core/indexes//multi.py:from pandas.core.arrays.categorical import factorize_from_iterables
pandas/core/indexes//multi.py:    indexer_from_factorized,
pandas/core/indexes//multi.py:        codes, levels = factorize_from_iterables(arrays)
pandas/core/indexes//multi.py:        codes, levels = factorize_from_iterables(iterables)
pandas/core/indexes//multi.py:        codes, uniques = algos.factorize(indexer, sort=True)
pandas/core/indexes//multi.py:            ok_codes, uniques = algos.factorize(indexer[mask], sort=True)
pandas/core/indexes//multi.py:            indexer = indexer_from_factorized(primary, primshp, compress=False)
Binary file pandas/core/indexes//__pycache__/multi.cpython-38.pyc matches

pls show where

Copy link
Contributor

Choose a reason for hiding this comment

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

hmm i guess its in the base class.

In any event. I don't think we actually want to change this.

Copy link
Member Author

Choose a reason for hiding this comment

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

Comment in existing tests seems to think this is already the behavior: https://github.com/pandas-dev/pandas/pull/33836/files#diff-8cf55ac38b6988b09a7f7f5d7280eb0fL360

Also should improve perf since this short-circuits the expensive part of factorize

if self.freq is not None:
# We must be unique, so can short-circuit (and retain freq)
codes = np.arange(len(self))
# TOOD: copy?
return codes, self[:]
return ExtensionArray.factorize(self, na_sentinel=na_sentinel)


DatetimeLikeArrayT = TypeVar("DatetimeLikeArrayT", bound="DatetimeLikeArrayMixin")

Expand Down
12 changes: 11 additions & 1 deletion pandas/core/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1402,7 +1402,17 @@ def memory_usage(self, deep=False):
),
)
def factorize(self, sort=False, na_sentinel=-1):
return algorithms.factorize(self, sort=sort, na_sentinel=na_sentinel)
codes, uniques = algorithms.factorize(
self._values, sort=sort, na_sentinel=na_sentinel
)
if isinstance(self, ABCIndexClass):
# use constructor instead of Index to get MultiIndex right
uniques = self._constructor(uniques)
else:
from pandas import Index

uniques = Index(uniques)
return codes, uniques

_shared_docs[
"searchsorted"
Expand Down
9 changes: 9 additions & 0 deletions pandas/core/indexes/datetimelike.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -602,6 +602,15 @@ def _shallow_copy(self, values=None, name: Label = lib.no_default):
result._cache = cache
return result

def factorize(self, sort=False, na_sentinel=-1):
if self.freq is not None and sort is False:
# we are unique, so can short-circuit, also can preserve freq
codes = np.arange(len(self))
return codes, self[:]
# TODO: In the sort=True case we could check for montonic_decreasing
# and operate on self[::-1]
return super().factorize(sort=sort, na_sentinel=na_sentinel)
Copy link
Member

Choose a reason for hiding this comment

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

This seems to be duplicating the array version, can't that be reused? (directly, or by having pd.factorize call the array version)

Copy link
Member

Choose a reason for hiding this comment

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

Can you check this comment?

Copy link
Member Author

Choose a reason for hiding this comment

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

Its close; the main difference is that the Index version returns an Index for uniques

Copy link
Member

Choose a reason for hiding this comment

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

But that is handled in factorize (depending on the input, it will wrap the uniques in an Index or not)?

Copy link
Member Author

Choose a reason for hiding this comment

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

im not clear on what youre suggesting? that we dont need to override this here at all?

Copy link
Member

Choose a reason for hiding this comment

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

Well, you might need to update factorize to preserve freq (it's using _shallow_copy right now)

(or we could decide that factorize is not an operation that should preserve the freq)

Copy link
Member Author

Choose a reason for hiding this comment

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

Well, you might need to update factorize to preserve freq (it's using _shallow_copy right now)

If we're getting into Index-subclass-specific logic, I think that belongs on the Index subclass.

Copy link
Member Author

Choose a reason for hiding this comment

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

(or we could decide that factorize is not an operation that should preserve the freq)

I'd be OK with this

Copy link
Member

Choose a reason for hiding this comment

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

If we're getting into Index-subclass-specific logic, I think that belongs on the Index subclass.

Not necessarily, I would say, as it already has index-specific handling as well.
But there is no "shallow_copy"-like method that preserves attributes like freq ?

Copy link
Member Author

Choose a reason for hiding this comment

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

But there is no "shallow_copy"-like method that preserves attributes like freq ?

There used to be, but it was being used in places where freq shouldnt be retained, so was changed as it was a footgun.


# --------------------------------------------------------------------
# Set Operation Methods

Expand Down
42 changes: 27 additions & 15 deletions pandas/tests/indexes/datetimes/test_datetime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -328,10 +328,12 @@ def test_factorize(self):
arr, idx = idx1.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, exp_idx)
assert idx.freq == exp_idx.freq

arr, idx = idx1.factorize(sort=True)
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, exp_idx)
assert idx.freq == exp_idx.freq

# tz must be preserved
idx1 = idx1.tz_localize("Asia/Tokyo")
Expand All @@ -340,6 +342,7 @@ def test_factorize(self):
arr, idx = idx1.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, exp_idx)
assert idx.freq == exp_idx.freq

idx2 = pd.DatetimeIndex(
["2014-03", "2014-03", "2014-02", "2014-01", "2014-03", "2014-01"]
Expand All @@ -350,49 +353,58 @@ def test_factorize(self):
arr, idx = idx2.factorize(sort=True)
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, exp_idx)
assert idx.freq == exp_idx.freq

exp_arr = np.array([0, 0, 1, 2, 0, 2], dtype=np.intp)
exp_idx = DatetimeIndex(["2014-03", "2014-02", "2014-01"])
arr, idx = idx2.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, exp_idx)
assert idx.freq == exp_idx.freq

# freq must be preserved
idx3 = date_range("2000-01", periods=4, freq="M", tz="Asia/Tokyo")
exp_arr = np.array([0, 1, 2, 3], dtype=np.intp)
arr, idx = idx3.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, idx3)
assert idx.freq == idx3.freq

def test_factorize_tz(self, tz_naive_fixture):
def test_factorize_tz(self, tz_naive_fixture, index_or_series):
tz = tz_naive_fixture
# GH#13750
base = pd.date_range("2016-11-05", freq="H", periods=100, tz=tz)
idx = base.repeat(5)

exp_arr = np.arange(100, dtype=np.intp).repeat(5)

for obj in [idx, pd.Series(idx)]:
arr, res = obj.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
expected = base._with_freq(None)
tm.assert_index_equal(res, expected)
obj = index_or_series(idx)

def test_factorize_dst(self):
arr, res = obj.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
expected = base._with_freq(None)
tm.assert_index_equal(res, expected)
assert res.freq == expected.freq

def test_factorize_dst(self, index_or_series):
# GH 13750
idx = pd.date_range("2016-11-06", freq="H", periods=12, tz="US/Eastern")
obj = index_or_series(idx)

for obj in [idx, pd.Series(idx)]:
arr, res = obj.factorize()
tm.assert_numpy_array_equal(arr, np.arange(12, dtype=np.intp))
tm.assert_index_equal(res, idx)
arr, res = obj.factorize()
tm.assert_numpy_array_equal(arr, np.arange(12, dtype=np.intp))
tm.assert_index_equal(res, idx)
if index_or_series is Index:
assert res.freq == idx.freq

idx = pd.date_range("2016-06-13", freq="H", periods=12, tz="US/Eastern")
obj = index_or_series(idx)

for obj in [idx, pd.Series(idx)]:
arr, res = obj.factorize()
tm.assert_numpy_array_equal(arr, np.arange(12, dtype=np.intp))
tm.assert_index_equal(res, idx)
arr, res = obj.factorize()
tm.assert_numpy_array_equal(arr, np.arange(12, dtype=np.intp))
tm.assert_index_equal(res, idx)
if index_or_series is Index:
assert res.freq == idx.freq

@pytest.mark.parametrize(
"arr, expected",
Expand Down
3 changes: 3 additions & 0 deletions pandas/tests/indexes/timedeltas/test_timedelta.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,17 +76,20 @@ def test_factorize(self):
arr, idx = idx1.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, exp_idx)
assert idx.freq == exp_idx.freq

arr, idx = idx1.factorize(sort=True)
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, exp_idx)
assert idx.freq == exp_idx.freq

# freq must be preserved
idx3 = timedelta_range("1 day", periods=4, freq="s")
exp_arr = np.array([0, 1, 2, 3], dtype=np.intp)
arr, idx = idx3.factorize()
tm.assert_numpy_array_equal(arr, exp_arr)
tm.assert_index_equal(idx, idx3)
assert idx.freq == idx3.freq

def test_sort_values(self):

Expand Down