Skip to content

Commit

Permalink
PERF: Groupby.shift dont re-call libgroupby.group_shift_indexer (pand…
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored and feefladder committed Sep 7, 2021
1 parent f50d447 commit 21ec89a
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 8 deletions.
18 changes: 11 additions & 7 deletions pandas/core/groupby/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -3031,15 +3031,19 @@ def shift(self, periods=1, freq=None, axis=0, fill_value=None):
if freq is not None or axis != 0:
return self.apply(lambda x: x.shift(periods, freq, axis, fill_value))

return self._get_cythonized_result(
"group_shift_indexer",
numeric_only=False,
cython_dtype=np.dtype(np.int64),
needs_ngroups=True,
result_is_index=True,
periods=periods,
ids, _, ngroups = self.grouper.group_info
res_indexer = np.zeros(len(ids), dtype=np.int64)

libgroupby.group_shift_indexer(res_indexer, ids, ngroups, periods)

obj = self._obj_with_exclusions

res = obj._reindex_with_indexers(
{self.axis: (obj.axes[self.axis], res_indexer)},
fill_value=fill_value,
allow_dups=True,
)
return res

@final
@Substitution(name="groupby")
Expand Down
2 changes: 1 addition & 1 deletion pandas/tests/groupby/transform/test_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,7 @@ def test_transform_axis_1(request, transformation_func, using_array_manager):
result = df.groupby([0, 0, 1], axis=1).transform(transformation_func, *args)
expected = df.T.groupby([0, 0, 1]).transform(transformation_func, *args).T

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

Expand Down

0 comments on commit 21ec89a

Please sign in to comment.