Brief Description

An auto_agg method where you just put the columns you want and it automatically aggregates the dataframe.

For example, I could do something like:

df.auto_agg(['state', 'age_mean'])

On a dataframe that has a state and age column and it will automatically get the average age by state.

Example API

Here's some basic code:

import pandas as pd

colnames = ('id', 'state', 'state_abbr', 'amount')

df = pd.DataFrame.from_records((
	('Person1', 'California', 'CA', 20),
	('Person2', 'California', 'CA', 30),
	('Person3', 'New York', 'NY', 15),
	('Person4', 'New York', 'NY', 10),
	('Person5', 'New York', 'NY', 45)), columns=colnames)


new_columns = ['state_abbr', 'state_count']


def flatten_dict(x):
    return {k: v for d in x for k, v in d.items()}


def auto_agg(df, new_columns):
    cols = set(df.columns)
    new_columns_set = set(new_columns)
    group_cols = cols.intersection(new_columns_set)
    agg_cols = new_columns_set.difference(cols)
    agg_groups = [{col: agg_col} for col in cols for agg_col in agg_cols if agg_col.startswith(col)]
    this_dict = dict()
    for i in agg_groups:
        value = list(i.values())[0]
        key = list(i.keys())[0]
        this_dict[value] = ''
        current_value = this_dict.get(value)
        if len(key) > len(current_value):
            this_dict[value] = key
    agg_groups_rename = dict([(value, key) for key, value in this_dict.items()])  
    agg_groups_agg = [{i[0]: i[1].split(i[0])[1:][0][1:]} for i in agg_groups_rename.items()]
    agg_groups_agg = flatten_dict(agg_groups_agg)
    group_cols_list = list(group_cols)
    return df.groupby(group_cols_list, as_index=False).agg(agg_groups_agg).rename(columns=agg_groups_rename)


auto_agg(df, new_columns)

There's a problem when column names are subsets of eachother, which is why I have an example with state_abbr and state. Sometimes the set math gets confused when trying to split the column names and the functions.

I think a basic workaround would be to be able to "force" the functions on specific columns without us needing to guess, like this:

new_columns = ['state', 'amount_sum', {'id': 'count'}]
auto_agg(df, new_columns)
# or df.auto_agg(new_columns)

If we say something like "if your choice of columns mess up, then be explicit" we might not need this part of the code, since we won't need to find the biggest or smallest version of the column name.

for i in agg_groups:
        value = list(i.values())[0]
        key = list(i.keys())[0]
        this_dict[value] = ''
        current_value = this_dict.get(value)
        if len(key) > len(current_value):
            this_dict[value] = key