Repository: pyjanitor-devs/pyjanitor

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!


Discussion:
Hi, the example API you mentioned would return grouped output, but it will only reveal so for the class that is on the top after sort_values.

The output is like this if groupby column is species and sort column is petal length. 

sepal_length | sepal_width | petal_length | petal_width | species
-- | -- | -- | -- | --
4.6 | 3.6 | 1.0 | 0.2 | setosa
4.3 | 3.0 | 1.1 | 0.1 | setosa
5.8 | 4.0 | 1.2 | 0.2 | setosa
5.0 | 3.2 | 1.2 | 0.2 | setosa
5.5 | 3.5 | 1.3 | 0.2 | setosa

However, `top protein mutants in each cluster` makes me believe that it is required to return the top few for each species, which can be achieved by 

```
def groupby_topk_2(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.sort_values([groupby_column_name,sort_column_name],**sort_values_kwargs).groupby(groupby_column_name).head(k)
```
whose output for the same would be 

sepal_length | sepal_width | petal_length | petal_width | species
-- | -- | -- | -- | --
4.6 | 3.6 | 1.0 | 0.2 | setosa
4.3 | 3.0 | 1.1 | 0.1 | setosa
5.8 | 4.0 | 1.2 | 0.2 | setosa
5.0 | 3.2 | 1.2 | 0.2 | setosa
4.7 | 3.2 | 1.3 | 0.2 | setosa
5.1 | 2.5 | 3.0 | 1.1 | versicolor
4.9 | 2.4 | 3.3 | 1.0 | versicolor
5.0 | 2.3 | 3.3 | 1.0 | versicolor
5.0 | 2.0 | 3.5 | 1.0 | versicolor
5.7 | 2.6 | 3.5 | 1.0 | versicolor
4.9 | 2.5 | 4.5 | 1.7 | virginica
6.2 | 2.8 | 4.8 | 1.8 | virginica
6.0 | 3.0 | 4.8 | 1.8 | virginica
5.6 | 2.8 | 4.9 | 2.0 | virginica
6.3 | 2.7 | 4.9 | 1.8 | virginica

`

Which of the two outputs is to be achieved?

Notebook that I used for reference and testing: https://colab.research.google.com/drive/1o9_MMXM6JEhfZI4a2c_ZSRyBvMyA_3Qn#scrollTo=fI-96IWoCYol&uniqifier=1


Sorry, notebook link I shared earlier was not correct.

Here is the correct link: https://colab.research.google.com/drive/1o9_MMXM6JEhfZI4a2c_ZSRyBvMyA_3Qn?usp=sharing
@mphirke that’s a great catch! In my original post, I actually put the `.head(k)` piece in the wrong place. I’ve edited the code block so it reflects the correct thing. Also confirmed that it works in your Colab notebook.

You seem interested in this functionality; would you like to work on a PR to contribute it? If so, I’m happy to guide you through the process!
I have not contributed to this repo before, but yes, I would like to work on it. 
Awesomeness! We have docs on how to get started contributing. Definitely head over to the [contributing page](https://pyjanitor.readthedocs.io/contributing.html) to see how to get setup and started. If you’re familiar with the `git` way of working (i.e. GitFlow), then you’ll definitely feel at home here!

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- 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
--- a/scripts/check-autodoc.py
+++ b/scripts/check-autodoc.py
@@ -70,5 +70,7 @@ def extract_folder_names(test_dir: Path) -> Tuple[Path, str]:
 
 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
--- /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})