Skip to content

Commit

Permalink
Fix return type of DataFrame.argsort (#7706)
Browse files Browse the repository at this point in the history
Fixes: #7577 

This PR fixes the return type of `DataFrame.argsort` to return `cudf.Series` instead of a `NumericalColumn`.

Authors:
  - GALI PREM SAGAR (@galipremsagar)

Approvers:
  - Keith Kraus (@kkraus14)

URL: #7706
  • Loading branch information
galipremsagar authored Mar 24, 2021
1 parent aa7ca46 commit df3c0f0
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 1 deletion.
24 changes: 23 additions & 1 deletion python/cudf/cudf/core/dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -3841,10 +3841,32 @@ def argsort(self, ascending=True, na_position="last"):
- Support axis='index' only.
- Not supporting: inplace, kind
- Ascending can be a list of bools to control per column
Examples
--------
>>> import cudf
>>> df = cudf.DataFrame({'a':[10, 0, 2], 'b':[-10, 10, 1]})
>>> df
a b
0 10 -10
1 0 10
2 2 1
>>> inds = df.argsort()
>>> inds
0 1
1 2
2 0
dtype: int32
>>> df.take(inds)
a b
1 0 10
2 2 1
0 10 -10
"""
return self._get_sorted_inds(
inds_col = self._get_sorted_inds(
ascending=ascending, na_position=na_position
)
return cudf.Series(inds_col)

@annotate("SORT_INDEX", color="red", domain="cudf_python")
def sort_index(
Expand Down
21 changes: 21 additions & 0 deletions python/cudf/cudf/tests/test_dataframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -8495,3 +8495,24 @@ def test_explode(data, labels, ignore_index, p_index, label_to_explode):
got = gdf.explode(label_to_explode, ignore_index)

assert_eq(expect, got, check_dtype=False)


@pytest.mark.parametrize(
"df,ascending,expected",
[
(
cudf.DataFrame({"a": [10, 0, 2], "b": [-10, 10, 1]}),
True,
cudf.Series([1, 2, 0], dtype="int32"),
),
(
cudf.DataFrame({"a": [10, 0, 2], "b": [-10, 10, 1]}),
False,
cudf.Series([0, 2, 1], dtype="int32"),
),
],
)
def test_dataframe_argsort(df, ascending, expected):
actual = df.argsort(ascending=ascending)

assert_eq(actual, expected)

0 comments on commit df3c0f0

Please sign in to comment.