diff --git a/python/cudf/cudf/core/series.py b/python/cudf/cudf/core/series.py index dcc6f0df164b..8d6f772e97c5 100644 --- a/python/cudf/cudf/core/series.py +++ b/python/cudf/cudf/core/series.py @@ -1154,6 +1154,59 @@ def to_frame(self, name: Hashable = no_default) -> DataFrame: self._propagate_metadata(res) return res + @_performance_tracking + def unstack(self, level=-1, fill_value=None, sort: bool = True): + """ + Unstack, also known as pivot, Series with MultiIndex to produce + DataFrame. + + Parameters + ---------- + level : int, str, or list of these, default last level + Level(s) to unstack, can pass level name. + fill_value + Non-functional argument provided for compatibility with Pandas. + sort : bool, default True + Sort the level(s) in the resulting MultiIndex columns. + + Returns + ------- + DataFrame + Unstacked Series. + + Examples + -------- + >>> import cudf + >>> s = cudf.Series( + ... [1, 2, 3, 4], + ... index=cudf.MultiIndex.from_product([["one", "two"], ["a", "b"]]), + ... ) + >>> s + one a 1 + b 2 + two a 3 + b 4 + dtype: int64 + >>> s.unstack(level=-1) + a b + one 1 2 + two 3 4 + """ + if not isinstance(self.index, cudf.MultiIndex): + raise ValueError( + "index must be a MultiIndex to unstack, " + f"{type(self.index)} was passed" + ) + result = self.to_frame().unstack( + level=level, fill_value=fill_value, sort=sort + ) + if result.columns.nlevels == 1: + # No level was actually unstacked (e.g. level=[]); pandas + # returns the original Series unchanged in that case. + return self.copy(deep=False) + result.columns = result.columns.droplevel(0) + return result + @_performance_tracking def memory_usage(self, index: bool = True, deep: bool = False) -> int: """ diff --git a/python/cudf/cudf/tests/reshape/test_unstack.py b/python/cudf/cudf/tests/reshape/test_unstack.py index 3a2462c4b49a..128f0da41a9e 100644 --- a/python/cudf/cudf/tests/reshape/test_unstack.py +++ b/python/cudf/cudf/tests/reshape/test_unstack.py @@ -105,3 +105,41 @@ def test_unstack_index_invalid(): ), ): gdf.unstack() + + +@pytest.mark.parametrize("level", [-1, 0, 1, "foo", "bar"]) +@pytest.mark.parametrize("name", [None, "quux"]) +def test_series_unstack_multiindex(level, name): + index = pd.MultiIndex.from_tuples( + [ + ("one", "a"), + ("one", "b"), + ("two", "a"), + ("two", "b"), + ], + names=["foo", "bar"], + ) + ps = pd.Series([1, 2, 3, 4], index=index, name=name) + gs = cudf.from_pandas(ps) + assert_eq(ps.unstack(level=level), gs.unstack(level=level)) + + +def test_series_unstack_index_invalid(): + gs = cudf.Series([1, 2, 3], index=["a", "b", "c"]) + with pytest.raises( + ValueError, + match=re.escape( + "index must be a MultiIndex to unstack, " + " was passed" + ), + ): + gs.unstack() + + +def test_series_unstack_empty_level_is_a_noop(): + index = pd.MultiIndex.from_tuples( + [("one", "a"), ("one", "b"), ("two", "a"), ("two", "b")] + ) + ps = pd.Series([1, 2, 3, 4], index=index, name="v") + gs = cudf.from_pandas(ps) + assert_eq(ps.unstack(level=[]), gs.unstack(level=[]))