Skip to content

Commit

Permalink
Fix stubname equal suffix in wide to long function
Browse files Browse the repository at this point in the history
  • Loading branch information
tqa236 committed Mar 19, 2024
1 parent 38086f1 commit ddad4a9
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 1 deletion.
8 changes: 7 additions & 1 deletion pandas/core/reshape/melt.py
Original file line number Diff line number Diff line change
Expand Up @@ -603,6 +603,9 @@ def get_var_names(df, stub: str, sep: str, suffix: str):
return df.columns[df.columns.str.match(regex)]

def melt_stub(df, stub: str, i, j, value_vars, sep: str):
# Ensure value_name and var_name are different when passing to melt
j_original = j
j = f"{j}_1" if stub == j else j
newdf = melt(
df,
id_vars=i,
Expand All @@ -619,7 +622,10 @@ def melt_stub(df, stub: str, i, j, value_vars, sep: str):
# TODO: anything else to catch?
pass

return newdf.set_index(i + [j])
newdf = newdf.set_index(i + [j])
if j != j_original:
newdf.index = newdf.index.set_names(j_original, level=-1)
return newdf

if not is_list_like(stubnames):
stubnames = [stubnames]
Expand Down
27 changes: 27 additions & 0 deletions pandas/tests/reshape/test_melt.py
Original file line number Diff line number Diff line change
Expand Up @@ -1218,6 +1218,33 @@ def test_missing_stubname(self, dtype):
expected.index = expected.index.set_levels(new_level, level=0)
tm.assert_frame_equal(result, expected)

@pytest.mark.parametrize("stubnames", ["year", ["year"]])
def test_stubname_equal_suffix(self, stubnames):
# https://github.com/pandas-dev/pandas/issues/46939
df = DataFrame(
{
"year1": {0: 4.5, 1: 1.7},
"year2": {0: 2.5, 1: 1.2},
"X": dict(zip(range(2), range(2, 4))),
}
)
df["id"] = df.index
result = wide_to_long(
df,
stubnames=stubnames,
i="id",
j="year",
)
expected = DataFrame(
[[2, 4.5], [3, 1.7], [2, 2.5], [3, 1.2]],
columns=["X", "year"],
index=pd.MultiIndex.from_arrays(
[[0, 1, 0, 1], [1, 1, 2, 2]],
names=["id", "year"],
),
)
tm.assert_frame_equal(result, expected)


def test_wide_to_long_pyarrow_string_columns():
# GH 57066
Expand Down

0 comments on commit ddad4a9

Please sign in to comment.