Skip to content

Commit d243ca0

Browse files
jbrockmendeljreback
authored andcommitted
REF: PeriodIndex.get_loc (#31021)
1 parent bea5789 commit d243ca0

File tree

3 files changed

+92
-53
lines changed

3 files changed

+92
-53
lines changed

pandas/_libs/index.pyx

+1-1
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,7 @@ cdef class TimedeltaEngine(DatetimeEngine):
498498
cdef class PeriodEngine(Int64Engine):
499499

500500
cdef _get_index_values(self):
501-
return super(PeriodEngine, self).vgetter()
501+
return super(PeriodEngine, self).vgetter().view("i8")
502502

503503
cdef void _call_map_locations(self, values):
504504
# super(...) pattern doesn't seem to work with `cdef`

pandas/core/indexes/period.py

+69-52
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import numpy as np
55

66
from pandas._libs import index as libindex
7-
from pandas._libs.tslibs import NaT, frequencies as libfrequencies, iNaT, resolution
7+
from pandas._libs.tslibs import NaT, frequencies as libfrequencies, resolution
88
from pandas._libs.tslibs.period import Period
99
from pandas.util._decorators import Appender, Substitution, cache_readonly
1010

@@ -17,6 +17,7 @@
1717
is_float_dtype,
1818
is_integer,
1919
is_integer_dtype,
20+
is_list_like,
2021
is_object_dtype,
2122
pandas_dtype,
2223
)
@@ -42,7 +43,6 @@
4243
)
4344
from pandas.core.indexes.datetimes import DatetimeIndex, Index
4445
from pandas.core.indexes.numeric import Int64Index
45-
from pandas.core.missing import isna
4646
from pandas.core.ops import get_op_result_name
4747
from pandas.core.tools.datetimes import DateParseError, parse_time_string
4848

@@ -507,42 +507,43 @@ def get_value(self, series, key):
507507
Fast lookup of value from 1-dimensional ndarray. Only use this if you
508508
know what you're doing
509509
"""
510-
s = com.values_from_object(series)
511-
try:
512-
value = super().get_value(s, key)
513-
except (KeyError, IndexError):
514-
if isinstance(key, str):
515-
asdt, parsed, reso = parse_time_string(key, self.freq)
516-
grp = resolution.Resolution.get_freq_group(reso)
517-
freqn = resolution.get_freq_group(self.freq)
518-
519-
vals = self._ndarray_values
520-
521-
# if our data is higher resolution than requested key, slice
522-
if grp < freqn:
523-
iv = Period(asdt, freq=(grp, 1))
524-
ord1 = iv.asfreq(self.freq, how="S").ordinal
525-
ord2 = iv.asfreq(self.freq, how="E").ordinal
526-
527-
if ord2 < vals[0] or ord1 > vals[-1]:
528-
raise KeyError(key)
529-
530-
pos = np.searchsorted(self._ndarray_values, [ord1, ord2])
531-
key = slice(pos[0], pos[1] + 1)
532-
return series[key]
533-
elif grp == freqn:
534-
key = Period(asdt, freq=self.freq).ordinal
535-
return com.maybe_box(
536-
self, self._int64index.get_value(s, key), series, key
537-
)
538-
else:
510+
if is_integer(key):
511+
return series.iat[key]
512+
513+
if isinstance(key, str):
514+
asdt, parsed, reso = parse_time_string(key, self.freq)
515+
grp = resolution.Resolution.get_freq_group(reso)
516+
freqn = resolution.get_freq_group(self.freq)
517+
518+
vals = self._ndarray_values
519+
520+
# if our data is higher resolution than requested key, slice
521+
if grp < freqn:
522+
iv = Period(asdt, freq=(grp, 1))
523+
ord1 = iv.asfreq(self.freq, how="S").ordinal
524+
ord2 = iv.asfreq(self.freq, how="E").ordinal
525+
526+
if ord2 < vals[0] or ord1 > vals[-1]:
539527
raise KeyError(key)
540528

541-
period = Period(key, self.freq)
542-
key = period.value if isna(period) else period.ordinal
543-
return com.maybe_box(self, self._int64index.get_value(s, key), series, key)
544-
else:
545-
return com.maybe_box(self, value, series, key)
529+
pos = np.searchsorted(self._ndarray_values, [ord1, ord2])
530+
key = slice(pos[0], pos[1] + 1)
531+
return series[key]
532+
elif grp == freqn:
533+
key = Period(asdt, freq=self.freq)
534+
loc = self.get_loc(key)
535+
return series.iloc[loc]
536+
else:
537+
raise KeyError(key)
538+
539+
elif isinstance(key, Period) or key is NaT:
540+
ordinal = key.ordinal if key is not NaT else NaT.value
541+
loc = self._engine.get_loc(ordinal)
542+
return series[loc]
543+
544+
# slice, PeriodIndex, np.ndarray, List[Period]
545+
value = Index.get_value(self, series, key)
546+
return com.maybe_box(self, value, series, key)
546547

547548
@Appender(_index_shared_docs["get_indexer"] % _index_doc_kwargs)
548549
def get_indexer(self, target, method=None, limit=None, tolerance=None):
@@ -579,36 +580,52 @@ def get_indexer_non_unique(self, target):
579580

580581
def get_loc(self, key, method=None, tolerance=None):
581582
"""
582-
Get integer location for requested label
583+
Get integer location for requested label.
584+
585+
Parameters
586+
----------
587+
key : Period, NaT, str, or datetime
588+
String or datetime key must be parseable as Period.
583589
584590
Returns
585591
-------
586-
loc : int
592+
loc : int or ndarray[int64]
593+
594+
Raises
595+
------
596+
KeyError
597+
Key is not present in the index.
598+
TypeError
599+
If key is listlike or otherwise not hashable.
587600
"""
588-
try:
589-
return self._engine.get_loc(key)
590-
except KeyError:
591-
if is_integer(key):
592-
raise
593601

602+
if isinstance(key, str):
594603
try:
595604
asdt, parsed, reso = parse_time_string(key, self.freq)
596605
key = asdt
597-
except TypeError:
598-
pass
599606
except DateParseError:
600607
# A string with invalid format
601608
raise KeyError(f"Cannot interpret '{key}' as period")
602609

603-
try:
604-
key = Period(key, freq=self.freq)
605-
except ValueError:
606-
# we cannot construct the Period
607-
# as we have an invalid type
608-
raise KeyError(key)
610+
elif is_integer(key):
611+
# Period constructor will cast to string, which we dont want
612+
raise KeyError(key)
613+
614+
try:
615+
key = Period(key, freq=self.freq)
616+
except ValueError:
617+
# we cannot construct the Period
618+
# as we have an invalid type
619+
if is_list_like(key):
620+
raise TypeError(f"'{key}' is an invalid key")
621+
raise KeyError(key)
622+
623+
ordinal = key.ordinal if key is not NaT else key.value
624+
try:
625+
return self._engine.get_loc(ordinal)
626+
except KeyError:
609627

610628
try:
611-
ordinal = iNaT if key is NaT else key.ordinal
612629
if tolerance is not None:
613630
tolerance = self._convert_tolerance(tolerance, np.asarray(key))
614631
return self._int64index.get_loc(ordinal, method, tolerance)

pandas/tests/indexes/period/test_indexing.py

+22
Original file line numberDiff line numberDiff line change
@@ -451,6 +451,28 @@ def test_get_loc(self):
451451
tm.assert_numpy_array_equal(idx2.get_loc(p2), expected_idx2_p2)
452452
tm.assert_numpy_array_equal(idx2.get_loc(str(p2)), expected_idx2_p2)
453453

454+
def test_get_loc_integer(self):
455+
dti = pd.date_range("2016-01-01", periods=3)
456+
pi = dti.to_period("D")
457+
with pytest.raises(KeyError, match="16801"):
458+
pi.get_loc(16801)
459+
460+
pi2 = dti.to_period("Y") # duplicates, ordinals are all 46
461+
with pytest.raises(KeyError, match="46"):
462+
pi2.get_loc(46)
463+
464+
def test_get_value_integer(self):
465+
dti = pd.date_range("2016-01-01", periods=3)
466+
pi = dti.to_period("D")
467+
ser = pd.Series(range(3), index=pi)
468+
with pytest.raises(IndexError, match="is out of bounds for axis 0 with size 3"):
469+
pi.get_value(ser, 16801)
470+
471+
pi2 = dti.to_period("Y") # duplicates, ordinals are all 46
472+
ser2 = pd.Series(range(3), index=pi2)
473+
with pytest.raises(IndexError, match="is out of bounds for axis 0 with size 3"):
474+
pi2.get_value(ser2, 46)
475+
454476
def test_is_monotonic_increasing(self):
455477
# GH 17717
456478
p0 = pd.Period("2017-09-01")

0 commit comments

Comments
 (0)