Repository: ericmjl/pyjanitor (python)

Issue:
[ENH]: Groupby Top K
# Brief Description

<!-- Please provide a brief description of what you'd like to propose. -->

I would like to propose a function that allows users to take in a dataframe and return another dataframe that has the top "k" values after grouping by another column.

This has come up multiple times at work, when I have to pick out "top protein mutants in each cluster".

Posting this here in case anybody wants to take a formal stab at it, otherwise I'll get to it later.

# Example API

```python
def groupby_topk(d: pd.DataFrame, groupby_column_name: Hashable, sort_column_name: Hashable, k: int, sort_values_kwargs: Dict) -> pd.DataFrame:
    """
    Return a dataframe that has the top "k" values per group.
    
    Some docstrings to flesh out:
    - Groups: defined by the "groupby" column.
    
    Taken inspiration from https://stackoverflow.com/a/27844045/1274908
    """
    return d.groupby(groupby_column_name).apply(lambda df: df.sort_values(sort_column_name, **sort_values_kwargs).head(k))
```

Give it a thumbs up if you're interested in seeing this too!

Pull request description:
[ENH] Added function groupby_topk
# PR Description

Please describe the changes proposed in the pull request:

- Creation of function groupby_topk that was suggested in https://github.com/ericmjl/pyjanitor/issues/710
- The function groups the dataframe by a specified column, then sorts it by another specified column and returns the top k entries.

**This PR resolves #710 **

<!-- 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. [x] 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.rst`.
<!-- We'd like to acknowledge your contributions! -->
3. [x] Add a line to `CHANGELOG.rst` 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.

## Quick Check

To do a very quick check that everything is correct, follow these steps below:

- [x] Run the command `make check` from pyjanitor's top-level directory. This will automatically run:
    - black formatting
    - flake8 checking
    - running the test suite
    - docs build

Once done, please check off the check-box above.

If `make check` does not work for you, you can execute the commands listed in the Makefile individually.

## Code Changes

<!-- If you have not made code changes, please feel free to delete this section. -->

If you are adding code changes, please ensure the following:

- [x] Ensure that you have added tests.
- [x] Run all tests (`$ pytest .`) locally on your machine.
    - [x] Check to ensure that test coverage covers the lines of code that you have added.
    - [ ] Ensure that all tests pass. Spark tests are failing on my machine despite no changes to them. I assume this is because I don't have Spark/JVM, rest all work. I will try again when if I can install Spark.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/AUTHORS.rst b/AUTHORS.rst
index ac1f3b7..b8b461b 100644
--- a/AUTHORS.rst
+++ b/AUTHORS.rst
@@ -90,3 +90,4 @@ Contributors
 - `@evan-anderson <https://github.com/evan-anderson>`_ | `contributions <https://github.com/ericmjl/pyjanitor/pulls?utf8=%E2%9C%93&q=is%3Aclosed+mentions%3evan-anderson>`_
 - `@smu095 <https://github.com/smu095>`_ | `contributions <https://github.com/ericmjl/pyjanitor/issues?q=is%3Aclosed+mentions%3smu095>`_
 - `@VPerrollaz <https://github.com/VPerrollaz>`_ | `contributions <https://github.com/ericmjl/pyjanitor/issues?q=is%3Aclosed+mentions%3AVPerrollaz>`_
+- `@mphirke <https://github.com/mphirke>`_ | `contributions <https://github.com/ericmjl/pyjanitor/issues?q=is%3Aclosed+mentions%3Amphirke>`_
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 43b9566..ffc84ca 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,6 +1,7 @@
 new version (on deck)
 =====================
 - [ENH] Updated groupby_agg function to account for null entries in the ``by`` argument. @samukweku
+- [ENH] Added function ``groupby_topk`` to janitor functions @mphirke
 
 
 v0.20.8
diff --git a/docs/reference/general_functions.rst b/docs/reference/general_functions.rst
index 0f19332..6afb00b 100644
--- a/docs/reference/general_functions.rst
+++ b/docs/reference/general_functions.rst
@@ -88,4 +88,5 @@ Other
     move
     toset
     unionize_dataframe_categories
+    groupby_topk
     
diff --git a/docs/reference/janitor.functions/janitor.groupby_topk.md b/docs/reference/janitor.functions/janitor.groupby_topk.md
new file mode 100644
index 0000000..b57fe7e
--- /dev/null
+++ b/docs/reference/janitor.functions/janitor.groupby_topk.md
@@ -0,0 +1,6 @@
+janitor.groupby\_topk
+====================
+
+.. currentmodule:: janitor
+
+.. autofunction:: groupby_topk
diff --git a/janitor/functions.py b/janitor/functions.py
index b4668a8..fb044ed 100644
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -1901,7 +1901,7 @@ def limit_column_characters(
     col_names = [col_name[:column_length] for col_name in col_names]
 
     col_name_set = set(col_names)
-    col_name_count = dict()
+    col_name_count = {}
 
     # If no columns are duplicates, we can skip the loops below.
     if len(col_name_set) == len(col_names):
@@ -3280,12 +3280,12 @@ class DataDescription:
     def __init__(self, data):
         """Initialize DataDescription class."""
         self._data = data
-        self._desc = dict()
+        self._desc = {}
 
     def _get_data_df(self) -> pd.DataFrame:
         df = self._data
 
-        data_dict = dict()
+        data_dict = {}
         data_dict["column_name"] = df.columns.tolist()
         data_dict["type"] = df.dtypes.tolist()
         data_dict["count"] = df.count().tolist()
@@ -4266,3 +4266,124 @@ def fill_direction(
                 df.loc[:, column].ffill(limit=limit).bfill(limit=limit)
             )
     return df
+
+
+@pf.register_dataframe_method
+def groupby_topk(
+    df: pd.DataFrame,
+    groupby_column_name: Hashable,
+    sort_column_name: Hashable,
+    k: int,
+    sort_values_kwargs: Dict = None,
+) -> 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`.
+
+    List of all sort_values() parameters can be found here_.
+
+    .. _here: https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.sort_values.html
+
+
+    .. code-block:: python
+
+        import pandas as pd
+        import janitor as jn
+
+        df = pd.DataFrame({'age' : [20, 22, 24, 23, 21, 22],
+                           'ID' : [1,2,3,4,5,6],
+                           'result' : ["pass", "fail", "pass",
+                                       "pass", "fail", "pass"]})
+
+        # Ascending top 3:
+        df.groupby_topk('result', 'age', 3)
+        #       age  ID  result
+        #result
+        #fail   21   5   fail
+        #       22   2   fail
+        #pass   20   1   pass
+        #       22   6   pass
+        #       23   4   pass
+
+        #Descending top 2:
+        df.groupby_topk('result', 'age', 2, {'ascending':False})
+        #       age  ID result
+        #result
+        #fail   22   2   fail
+        #       21   5   fail
+        #pass   24   3   pass
+        #       23   4   pass
+
+    Functional usage syntax:
+
+    .. code-block:: python
+
+        import pandas as pd
+        import janitor as jn
+
+        df = pd.DataFrame(...)
+        df = jn.groupby_topk(
+            df = df,
+            groupby_column_name = 'groupby_column',
+            sort_column_name = 'sort_column',
+            k = 5
+            )
+
+    Method-chaining usage syntax:
+
+    .. code-block:: python
+
+        import pandas as pd
+        import janitor as jn
+
+        df = (
+            pd.DataFrame(...)
+            .groupby_topk(
+            df = df,
+            groupby_column_name = 'groupby_column',
+            sort_column_name = 'sort_column',
+            k = 5
+            )
+        )
+
+    :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`.
+    :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 {}
+
+    # Check if groupby_column_name and sort_column_name exists in the dataframe
+    check_column(df, [groupby_column_name, sort_column_name])
+
+    # 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."
+        )
+
+    # 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`.")
+
+    return df.groupby(groupby_column_name).apply(
+        lambda d: d.sort_values(sort_column_name, **sort_values_kwargs).head(k)
+    )
diff --git a/scripts/check-autodoc.py b/scripts/check-autodoc.py
index e11899d..f514a6f 100644
--- a/scripts/check-autodoc.py
+++ b/scripts/check-autodoc.py
@@ -70,5 +70,7 @@ missing_funcs = set(function_names).difference(doc_functions)
 
 if len(missing_funcs) > 0:
     raise Exception(
-        f"The following functions have not yet been added: {missing_funcs}"
+        f"""The following functions have not yet been added: {missing_funcs}.
+         Please add the function names to docs/reference/general_functions.rst
+         under an appropriate section."""
     )


Test patch:
diff --git a/tests/functions/test_groupby_topk.py b/tests/functions/test_groupby_topk.py
new file mode 100644
index 0000000..55e9bb7
--- /dev/null
+++ b/tests/functions/test_groupby_topk.py
@@ -0,0 +1,61 @@
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+
+
+@pytest.fixture
+def df():
+    return pd.DataFrame(
+        [
+            {"age": 22, "major": "science", "ID": 145, "result": "pass"},
+            {"age": 20, "major": "history", "ID": 322, "result": "fail"},
+            {"age": 23, "major": "history", "ID": 321, "result": "pass"},
+            {"age": 21, "major": "law", "ID": 49, "result": "fail"},
+            {"age": 19, "major": "mathematics", "ID": 224, "result": "pass"},
+            {"age": 20, "major": "science", "ID": 132, "result": "pass"},
+        ]
+    )
+
+
+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)
+    )
+    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)
+    )
+    assert_frame_equal(
+        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):
+        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):
+        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):
+        df.groupby_topk("result", "age", -2)
+    with pytest.raises(ValueError):
+        df.groupby_topk("result", "age", 0)
+
+
+def test_inplace(df):
+    """Raise Key Error if inplace is True in sort_values_kwargs"""
+    with pytest.raises(KeyError):
+        df.groupby_topk("result", "age", 1, {"inplace": True})