Skip to content

Commit

Permalink
REG: DataFrame.shift with axis=1 and CategoricalIndex columns (#38504)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel committed Dec 17, 2020
1 parent d4b6233 commit d08f12c
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 4 deletions.
1 change: 0 additions & 1 deletion doc/source/whatsnew/v1.3.0.rst
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,6 @@ Numeric
^^^^^^^
- Bug in :meth:`DataFrame.quantile`, :meth:`DataFrame.sort_values` causing incorrect subsequent indexing behavior (:issue:`38351`)
- Bug in :meth:`DataFrame.select_dtypes` with ``include=np.number`` now retains numeric ``ExtensionDtype`` columns (:issue:`35340`)
-

Conversion
^^^^^^^^^^
Expand Down
7 changes: 5 additions & 2 deletions pandas/core/frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -4586,20 +4586,23 @@ def shift(
if axis == 1 and periods != 0 and fill_value is lib.no_default and ncols > 0:
# We will infer fill_value to match the closest column

# Use a column that we know is valid for our column's dtype GH#38434
label = self.columns[0]

if periods > 0:
result = self.iloc[:, :-periods]
for col in range(min(ncols, abs(periods))):
# TODO(EA2D): doing this in a loop unnecessary with 2D EAs
# Define filler inside loop so we get a copy
filler = self.iloc[:, 0].shift(len(self))
result.insert(0, col, filler, allow_duplicates=True)
result.insert(0, label, filler, allow_duplicates=True)
else:
result = self.iloc[:, -periods:]
for col in range(min(ncols, abs(periods))):
# Define filler inside loop so we get a copy
filler = self.iloc[:, -1].shift(len(self))
result.insert(
len(result.columns), col, filler, allow_duplicates=True
len(result.columns), label, filler, allow_duplicates=True
)

result.columns = self.columns.copy()
Expand Down
24 changes: 23 additions & 1 deletion pandas/tests/frame/methods/test_shift.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import pytest

import pandas as pd
from pandas import DataFrame, Index, Series, date_range, offsets
from pandas import CategoricalIndex, DataFrame, Index, Series, date_range, offsets
import pandas._testing as tm


Expand Down Expand Up @@ -292,3 +292,25 @@ def test_shift_dt64values_int_fill_deprecated(self):

expected = DataFrame({"A": [pd.Timestamp(0), pd.Timestamp(0)], "B": df2["A"]})
tm.assert_frame_equal(result, expected)

def test_shift_axis1_categorical_columns(self):
# GH#38434
ci = CategoricalIndex(["a", "b", "c"])
df = DataFrame(
{"a": [1, 3], "b": [2, 4], "c": [5, 6]}, index=ci[:-1], columns=ci
)
result = df.shift(axis=1)

expected = DataFrame(
{"a": [np.nan, np.nan], "b": [1, 3], "c": [2, 4]}, index=ci[:-1], columns=ci
)
tm.assert_frame_equal(result, expected)

# periods != 1
result = df.shift(2, axis=1)
expected = DataFrame(
{"a": [np.nan, np.nan], "b": [np.nan, np.nan], "c": [1, 3]},
index=ci[:-1],
columns=ci,
)
tm.assert_frame_equal(result, expected)

0 comments on commit d08f12c

Please sign in to comment.