From 7ec38f92b846384e4cbd8022f0c5d99c17b3f42f Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 4 Sep 2023 15:42:11 +0100 Subject: [PATCH] Add test of partial multiindex sorting --- python/cudf/cudf/tests/test_multiindex.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/python/cudf/cudf/tests/test_multiindex.py b/python/cudf/cudf/tests/test_multiindex.py index eedc9b0c174..56bd7d709b7 100644 --- a/python/cudf/cudf/tests/test_multiindex.py +++ b/python/cudf/cudf/tests/test_multiindex.py @@ -1897,3 +1897,26 @@ def test_multiindex_empty_slice_pandas_compatibility(): with cudf.option_context("mode.pandas_compatible", True): actual = cudf.from_pandas(expected) assert_eq(expected, actual, exact=False) + + +@pytest.mark.parametrize( + "levels", + itertools.chain.from_iterable( + itertools.permutations(range(3), n) for n in range(1, 4) + ), + ids=str, +) +def test_multiindex_sort_index_partial(levels): + df = pd.DataFrame( + { + "a": [3, 3, 3, 1, 1, 1, 2, 2], + "b": [4, 2, 7, -1, 11, -2, 7, 7], + "c": [4, 4, 2, 3, 3, 3, 1, 1], + "val": [1, 2, 3, 4, 5, 6, 7, 8], + } + ).set_index(["a", "b", "c"]) + cdf = cudf.from_pandas(df) + + expect = df.sort_index(level=levels, sort_remaining=True) + got = cdf.sort_index(level=levels, sort_remaining=True) + assert_eq(expect, got)