Skip to content

Commit 21ec89a

Browse files
jbrockmendelfeefladder
authored andcommitted
PERF: Groupby.shift dont re-call libgroupby.group_shift_indexer (pandas-dev#42885)
1 parent f50d447 commit 21ec89a

File tree

2 files changed

+12
-8
lines changed

2 files changed

+12
-8
lines changed

pandas/core/groupby/groupby.py

+11-7
Original file line numberDiff line numberDiff line change
@@ -3031,15 +3031,19 @@ def shift(self, periods=1, freq=None, axis=0, fill_value=None):
30313031
if freq is not None or axis != 0:
30323032
return self.apply(lambda x: x.shift(periods, freq, axis, fill_value))
30333033

3034-
return self._get_cythonized_result(
3035-
"group_shift_indexer",
3036-
numeric_only=False,
3037-
cython_dtype=np.dtype(np.int64),
3038-
needs_ngroups=True,
3039-
result_is_index=True,
3040-
periods=periods,
3034+
ids, _, ngroups = self.grouper.group_info
3035+
res_indexer = np.zeros(len(ids), dtype=np.int64)
3036+
3037+
libgroupby.group_shift_indexer(res_indexer, ids, ngroups, periods)
3038+
3039+
obj = self._obj_with_exclusions
3040+
3041+
res = obj._reindex_with_indexers(
3042+
{self.axis: (obj.axes[self.axis], res_indexer)},
30413043
fill_value=fill_value,
3044+
allow_dups=True,
30423045
)
3046+
return res
30433047

30443048
@final
30453049
@Substitution(name="groupby")

pandas/tests/groupby/transform/test_transform.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ def test_transform_axis_1(request, transformation_func, using_array_manager):
183183
result = df.groupby([0, 0, 1], axis=1).transform(transformation_func, *args)
184184
expected = df.T.groupby([0, 0, 1]).transform(transformation_func, *args).T
185185

186-
if transformation_func == "diff":
186+
if transformation_func in ["diff", "shift"]:
187187
# Result contains nans, so transpose coerces to float
188188
expected["b"] = expected["b"].astype("int64")
189189

0 commit comments

Comments
 (0)