Repository: pyjanitor-devs/pyjanitor (python)

Issue:
groupby_topk performance 
it is possible to use `nlargest/nsmallest` instead of sort_values for `groupby_topk`, this should offer more performance by avoiding a full sort (which is what nlargest/nsmallest achieves)

Pull request:
[ENH] faster groupby top k
# PR Description

Please describe the changes proposed in the pull request:

- Use nlargest/nsmallest instead ... avoids sorting the entire dataframe
- change column names to align more with pandas description and other existing functions (open to suggestions, or possible revert to existing names)
- deprecate sort_kwargs, since there is no use of `sort_values`
-  add `ascending` parameter to match existing parameter in the deprecated `sort_kwargs`

<!-- Doing so provides maintainers with context on what the PR is, and can help us more effectively review your PR. -->

<!-- Please also identify below which issue that has been raised that you are going to close. -->

**This PR resolves #1093 .**


# performance comparison ... view with a pinch of salt, as always

```py
# adapted from  https://stackoverflow.com/a/36702926/7175713
np.random.seed(3)
df = pd.DataFrame({
   'dep': np.random.choice(list('ABC'), 2_000_000),
  'manager_id': np.random.randint(0, 200_000, 2_000_000),
   'salary': np.random.randint(5000, 5006, 2_000_000)
 })

df.head()
  dep  manager_id  salary
0   C       12586    5002
1   A      153719    5004
2   B      178720    5004
3   A      162938    5001
4   A      122360    5000

# --------------------> PR

df.groupby_topk('dep', 'salary', 5)
   dep  manager_id  salary
0    C       44168    5000
1    C      135457    5000
2    C       42892    5000
3    C      130785    5000
4    C      155488    5000
5    A      122360    5000
6    A      136722    5000
7    A       14219    5000
8    A      162417    5000
9    A      115906    5000
10   B       66927    5000
11   B       79314    5000
12   B      150702    5000
13   B      186580    5000
14   B       34819    5000

df.groupby_topk('dep', 'salary', 5, ignore_index = False) 
    dep  manager_id  salary
11    C       44168    5000
13    C      135457    5000
42    C       42892    5000
73    C      130785    5000
118   C      155488    5000
4     A      122360    5000
14    A      136722    5000
30    A       14219    5000
46    A      162417    5000
49    A      115906    5000
29    B       66927    5000
33    B       79314    5000
38    B      150702    5000
51    B      186580    5000
76    B       34819    5000

%timeit df.groupby_topk('dep', 'salary', 5)
162 ms ± 12.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit df.groupby_topk('dep', 'salary', 5, ignore_index = False)
154 ms ± 1.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)



# ---------------------------> dev

df.groupby_topk('dep', 'salary', 5)
            dep  manager_id  salary
dep                                
A   1999996   A      131800    5000
    548719    A       57184    5000
    548720    A       59279    5000
    1585713   A       87340    5000
    548727    A      196661    5000
B   716069    B       89977    5000
    306266    B       34652    5000
    743267    B      125335    5000
    1877756   B      166187    5000
    1703182   B      101396    5000
C   1999997   C      180787    5000
    292147    C      105282    5000
    1609388   C       52935    5000
    1609401   C       55700    5000
    292136    C      111160    5000

df.groupby_topk('dep', 'salary', 5).droplevel(0) 
        dep  manager_id  salary
1999996   A      131800    5000
548719    A       57184    5000
548720    A       59279    5000
1585713   A       87340    5000
548727    A      196661    5000
716069    B       89977    5000
306266    B       34652    5000
743267    B      125335    5000
1877756   B      166187    5000
1703182   B      101396    5000
1999997   C      180787    5000
292147    C      105282    5000
1609388   C       52935    5000
1609401   C       55700    5000
292136    C      111160    5000

%timeit df.groupby_topk('dep', 'salary', 5)
297 ms ± 4.18 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
 ```


<!-- As you go down the PR template, please feel free to delete sections that are irrelevant. -->

# PR Checklist

<!-- This checklist exists for newcomers who are not yet familiar with our requirements. If you are experienced with
the project, please feel free to delete this section. -->

Please ensure that you have done the following:

1. [ ] PR in from a fork off your branch. Do not PR from `<your_username>`:`dev`, but rather from `<your_username>`:`<feature-branch_name>`.
<!-- Doing this helps us keep the commit history much cleaner than it would otherwise be. -->
2. [x] If you're not on the contributors list, add yourself to `AUTHORS.md`.
<!-- We'd like to acknowledge your contributions! -->
3. [x] Add a line to `CHANGELOG.md` under the latest version header (i.e. the one that is "on deck") describing the contribution.
    - Do use some discretion here; if there are multiple PRs that are related, keep them in a single line.

# Automatic checks

There will be automatic checks run on the PR. These include:

- Building a preview of the docs on Netlify
- Automatically linting the code
- Making sure the code is documented
- Making sure that all tests are passed
- Making sure that code coverage doesn't go down.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 10e21d5..c14e982 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@
 -   [ENH] Allow column selection/renaming within conditional_join. #1102 @samukweku.
 -   [ENH] New decorator `deprecated_kwargs` for breaking API. #1103 @Zeroto521
 -   [ENH] Extend select_columns to support non-string columns. #1105 @samukweku
+-   [ENH] Performance improvement for groupby_topk. #1093 @samukweku
 
 ## [v0.23.1] - 2022-05-03
 
diff --git a/janitor/functions/groupby_topk.py b/janitor/functions/groupby_topk.py
index abc44f1..e49050f 100644
--- a/janitor/functions/groupby_topk.py
+++ b/janitor/functions/groupby_topk.py
@@ -1,39 +1,46 @@
 """Implementation of the `groupby_topk` function"""
-from typing import Dict, Hashable
+from typing import Hashable, Union
 import pandas_flavor as pf
 import pandas as pd
 
 from janitor.utils import check_column
+from janitor.utils import check, deprecated_alias
 
 
 @pf.register_dataframe_method
+@deprecated_alias(groupby_column_name="by", sort_column_name="column")
 def groupby_topk(
     df: pd.DataFrame,
-    groupby_column_name: Hashable,
-    sort_column_name: Hashable,
+    by: Union[list, Hashable],
+    column: Hashable,
     k: int,
-    sort_values_kwargs: Dict = None,
+    dropna: bool = True,
+    ascending: bool = True,
+    ignore_index: bool = True,
 ) -> pd.DataFrame:
     """
     Return top `k` rows from a groupby of a set of columns.
 
-    Returns a DataFrame that has the top `k` values grouped by `groupby_column_name`
-    and sorted by `sort_column_name`.
-    Additional parameters to the sorting (such as `ascending=True`)
-    can be passed using `sort_values_kwargs`.
+    Returns a DataFrame that has the top `k` values per `column`,
+    grouped by `by`. Under the hood it uses `nlargest/nsmallest`,
+    for numeric columns, which avoids sorting the entire dataframe,
+    and is usually more performant. For non-numeric columns, `pd.sort_values`
+    is used.
+    No sorting is done to the `by` column(s); the order is maintained
+    in the final output.
 
-    List of all sort_values() parameters can be found
-    [here](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.sort_values.html).
 
     Example:
 
         >>> import pandas as pd
         >>> import janitor
-        >>> df = pd.DataFrame({
-        ...     "age": [20, 23, 22, 43, 21],
-        ...     "id": [1, 4, 6, 2, 5],
-        ...     "result": ["pass", "pass", "fail", "pass", "fail"]
-        ... })
+        >>> df = pd.DataFrame(
+        ...     {
+        ...         "age": [20, 23, 22, 43, 21],
+        ...         "id": [1, 4, 6, 2, 5],
+        ...         "result": ["pass", "pass", "fail", "pass", "fail"],
+        ...     }
+        ... )
         >>> df
            age  id result
         0   20   1   pass
@@ -41,59 +48,77 @@ def groupby_topk(
         2   22   6   fail
         3   43   2   pass
         4   21   5   fail
-        >>> df.groupby_topk('result', 'age', 3) # Ascending top 3
-        ... # doctest: +NORMALIZE_WHITESPACE
-                  age  id result
-        result
-        fail   4   21   5   fail
-               2   22   6   fail
-        pass   0   20   1   pass
-               1   23   4   pass
-               3   43   2   pass
-        >>> df.groupby_topk('result', 'age', 2, {'ascending':False}) # Descending top 2
-        ... # doctest: +NORMALIZE_WHITESPACE
-                  age  id result
-        result
-        fail   2   22   6   fail
-               4   21   5   fail
-        pass   3   43   2   pass
-               1   23   4   pass
+
+    Ascending top 3:
+
+        >>> df.groupby_topk(by="result", column="age", k=3)
+           age  id result
+        0   20   1   pass
+        1   23   4   pass
+        2   43   2   pass
+        3   21   5   fail
+        4   22   6   fail
+
+    Descending top 2:
+
+        >>> df.groupby_topk(
+        ...     by="result", column="age", k=2, ascending=False, ignore_index=False
+        ... )
+           age  id result
+        3   43   2   pass
+        1   23   4   pass
+        2   22   6   fail
+        4   21   5   fail
 
 
     :param df: A pandas DataFrame.
-    :param groupby_column_name: Column name to group input DataFrame `df` by.
-    :param sort_column_name: Name of the column to sort along the
-        input DataFrame `df`.
-    :param k: Number of top rows to return from each group after sorting.
-    :param sort_values_kwargs: Arguments to be passed to sort_values function.
-    :returns: A pandas DataFrame with top `k` rows that are grouped by
-        `groupby_column_name` column with each group sorted along the
-        column `sort_column_name`.
+    :param by: Column name(s) to group input DataFrame `df` by.
+    :param column: Name of the column that determines `k` rows
+        to return.
+    :param k: Number of top rows to return for each group.
+    :param dropna: If `True`, and `NA` values exist in `by`, the `NA`
+        values are not used in the groupby computation to get the relevant
+        `k` rows. If `False`, and `NA` values exist in `by`, then the `NA`
+        values are used in the groupby computation to get the relevant
+        `k` rows. The default is `True`.
+    :param ascending: Default is `True`. If `True`, the smallest top `k` rows,
+        determined by `column` are returned; if `False, the largest top `k` rows,
+        determined by `column` are returned.
+    :param ignore_index: Default `True`. If `True`,
+        the original index is ignored. If `False`, the original index
+        for the top `k` rows is retained.
+    :returns: A pandas DataFrame with top `k` rows per `column`, grouped by `by`.
     :raises ValueError: if `k` is less than 1.
-    :raises ValueError: if `groupby_column_name` not in DataFrame `df`.
-    :raises ValueError: if `sort_column_name` not in DataFrame `df`.
-    :raises KeyError: if `inplace:True` is present in `sort_values_kwargs`.
     """  # noqa: E501
 
-    # Convert the default sort_values_kwargs from None to empty Dict
-    sort_values_kwargs = sort_values_kwargs or {}
+    if isinstance(by, Hashable):
+        by = [by]
 
-    # Check if groupby_column_name and sort_column_name exists in the DataFrame
-    check_column(df, [groupby_column_name, sort_column_name])
+    check("by", by, [Hashable, list])
+
+    check_column(df, [column])
+    check_column(df, by)
 
-    # Check if k is greater than 0.
     if k < 1:
         raise ValueError(
-            "Numbers of rows per group to be returned must be greater than 0."
+            "Numbers of rows per group "
+            "to be returned must be greater than 0."
         )
 
-    # Check if inplace:True in sort values kwargs because it returns None
-    if (
-        "inplace" in sort_values_kwargs.keys()
-        and sort_values_kwargs["inplace"]
-    ):
-        raise KeyError("Cannot use `inplace=True` in `sort_values_kwargs`.")
+    indices = df.groupby(by=by, dropna=dropna, sort=False, observed=True)
+    indices = indices[column]
+
+    try:
+        if ascending:
+            indices = indices.nsmallest(n=k)
+        else:
+            indices = indices.nlargest(n=k)
+    except TypeError:
+        indices = indices.apply(
+            lambda d: d.sort_values(ascending=ascending).head(k)
+        )
 
-    return df.groupby(groupby_column_name).apply(
-        lambda d: d.sort_values(sort_column_name, **sort_values_kwargs).head(k)
-    )
+    indices = indices.index.get_level_values(-1)
+    if ignore_index:
+        return df.loc[indices].reset_index(drop=True)
+    return df.loc[indices]


Test patch:
diff --git a/tests/functions/test_groupby_topk.py b/tests/functions/test_groupby_topk.py
index 55e9bb7..b1bbfe9 100644
--- a/tests/functions/test_groupby_topk.py
+++ b/tests/functions/test_groupby_topk.py
@@ -5,6 +5,7 @@ from pandas.testing import assert_frame_equal
 
 @pytest.fixture
 def df():
+    """fixture for groupby_topk"""
     return pd.DataFrame(
         [
             {"age": 22, "major": "science", "ID": 145, "result": "pass"},
@@ -17,44 +18,75 @@ def df():
     )
 
 
+def test_dtype_by(df):
+    """Check dtype for by."""
+    with pytest.raises(TypeError):
+        df.groupby_topk(by={"result"}, column="age", k=2)
+
+
 def test_ascending_groupby_k_2(df):
     """Test ascending group by, k=2"""
-    expected = df.groupby("result").apply(
-        lambda d: d.sort_values("age").head(2)
+    expected = (
+        df.groupby("result", sort=False)
+        .apply(lambda d: d.sort_values("age").head(2))
+        .droplevel(0)
+    )
+    assert_frame_equal(
+        df.groupby_topk("result", "age", 2, ignore_index=False), expected
+    )
+
+
+def test_ascending_groupby_non_numeric(df):
+    """Test output for non-numeric column"""
+    expected = (
+        df.groupby("result", sort=False)
+        .apply(lambda d: d.sort_values("major").head(2))
+        .droplevel(0)
+    )
+    assert_frame_equal(
+        df.groupby_topk("result", "major", 2, ignore_index=False), expected
     )
-    assert_frame_equal(df.groupby_topk("result", "age", 2), expected)
 
 
 def test_descending_groupby_k_3(df):
     """Test descending group by, k=3"""
-    expected = df.groupby("result").apply(
-        lambda d: d.sort_values("age", ascending=False).head(3)
+    expected = (
+        df.groupby("result", sort=False)
+        .apply(lambda d: d.sort_values("age", ascending=False).head(3))
+        .droplevel(0)
+        .reset_index(drop=True)
     )
     assert_frame_equal(
-        df.groupby_topk("result", "age", 3, {"ascending": False}), expected
+        df.groupby_topk("result", "age", 3, ascending=False), expected
     )
 
 
 def test_wrong_groupby_column_name(df):
     """Raise Value Error if wrong groupby column name is provided."""
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError, match="RESULT not present in dataframe columns!"
+    ):
         df.groupby_topk("RESULT", "age", 3)
 
 
 def test_wrong_sort_column_name(df):
     """Raise Value Error if wrong sort column name is provided."""
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError, match="Age not present in dataframe columns!"
+    ):
         df.groupby_topk("result", "Age", 3)
 
 
 def test_negative_k(df):
     """Raises Value Error if k is less than 1 (negative or 0)."""
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError,
+        match="Numbers of rows per group.+",
+    ):
         df.groupby_topk("result", "age", -2)
-    with pytest.raises(ValueError):
-        df.groupby_topk("result", "age", 0)
 
 
+@pytest.mark.xfail(reason="sort_value_kwargs parameter deprecated.")
 def test_inplace(df):
     """Raise Key Error if inplace is True in sort_values_kwargs"""
     with pytest.raises(KeyError):