Skip to content

[ENH] Method for adding functionality to pandas groupby #1462

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 16 commits into
base: dev
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@

## [Unreleased]
- [ENH] Added `row_count` parameter for janitor.conditional_join - Issue #1269 @samukweku
- [ENG] Reverse deprecation of `pivot_wider()` -- Issue #1464
- [ENH] Reverse deprecation of `pivot_wider()` -- Issue #1464
- [ENH] Add accessor and method for pandas DataFrameGroupBy objects. - Issue #587 @samukweku

## [v0.31.0] - 2025-03-07

- [ENH] Added support for pd.Series.select - Issue #1394 @samukweku
Expand Down
22 changes: 20 additions & 2 deletions janitor/functions/select.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,17 +327,18 @@ def select_rows(
return _select(df, rows=list(args), invert=invert)


@pf.register_groupby_method
@pf.register_dataframe_method
@pf.register_series_method
@deprecated_alias(rows="index")
def select(
df: pd.DataFrame | pd.Series,
df: pd.DataFrame | pd.Series | DataFrameGroupBy,
*args: tuple,
index: Any = None,
columns: Any = None,
axis: str = "columns",
invert: bool = False,
) -> pd.DataFrame | pd.Series:
) -> pd.DataFrame | pd.Series | DataFrameGroupBy:
"""Method-chainable selection of rows and/or columns.

It accepts a string, shell-like glob strings `(*string*)`,
Expand Down Expand Up @@ -371,6 +372,8 @@ def select(
- `rows` keyword deprecated in favour of `index`.
- 0.31.0
- Add support for pd.Series.
- 0.32.0
- Add support for DataFrameGroupBy.

Examples:
>>> import pandas as pd
Expand Down Expand Up @@ -436,6 +439,10 @@ def select(
Returns:
A pandas DataFrame or Series with the specified rows and/or columns selected.
""" # noqa: E501
if args and isinstance(df, DataFrameGroupBy):
return get_columns(group=df, label=list(args))
if isinstance(df, DataFrameGroupBy):
return get_columns(group=df, label=[columns])
if args:
check("invert", invert, [bool])
if (index is not None) or (columns is not None):
Expand Down Expand Up @@ -478,6 +485,12 @@ def get_index_labels(
return index[_select_index(arg, df, axis)]


@refactored_function(
message=(
"This function will be deprecated in a 1.x release. "
"Please use `jn.select` instead."
)
)
def get_columns(
group: DataFrameGroupBy | SeriesGroupBy, label: Any
) -> DataFrameGroupBy | SeriesGroupBy:
Expand All @@ -488,6 +501,11 @@ def get_columns(

!!! info "New in version 0.25.0"

!!!note

This function will be deprecated in a 1.x release.
Please use `jn.select` instead.

Args:
group: A Pandas GroupBy object.
label: column(s) to select.
Expand Down
14 changes: 14 additions & 0 deletions tests/functions/test_select_columns.py
Original file line number Diff line number Diff line change
Expand Up @@ -479,6 +479,20 @@ def test_select_groupby(dataframe):
assert_frame_equal(expected, actual)


def test_select_groupby_args(dataframe):
"""Test output on a grouped object"""
expected = dataframe.select_dtypes("number").groupby(dataframe["a"]).sum()
actual = dataframe.groupby("a").select(is_numeric_dtype).sum()
assert_frame_equal(expected, actual)


def test_select_groupby_columns(dataframe):
"""Test output on a grouped object"""
expected = dataframe.select_dtypes("number").groupby(dataframe["a"]).sum()
actual = dataframe.groupby("a").select(columns=is_numeric_dtype).sum()
assert_frame_equal(expected, actual)


def test_select_str_multiindex(multiindex):
"""Test str selection on a MultiIndex - exact match"""
expected = multiindex.select_columns("bar")
Expand Down
Loading