How to refactor a lot of nested for loops for pandas filtering

user_10468726 0
Viewed 33

I am new to working with pandas and have the following (pseudo codish) function as part of a python program to create reports

def loop_over_dataframe(df: DataFrame, report_types: dict)
...
for key, value in report_types:
...
    filtered_df = df.query(value["filter"])
    for date in filtered_df["Date"].unique():
        date_df = filtered_df[fildered_df["Date"] == date]
        ...
        for event in date_df["Events"].unique():
            event_df = date_df[date_df[] == event]
            ...
            for company in event_df["CompanyCode"].unique():
                company_df = event_df[event_df["CompanyCode"] == company]
                    generate_report(company_df, key)

In a Code Review this came up as having a lot of indentations for one function. But I am basically just looping over the DataFrame for all reports that should be created. I don't see a point to split this over multiple functions, or should I? Is there maybe also a nicer way to do this in pandas other than for-looping all over the place?

pandas refactoring python
1 Answers

You can make the code much more concise by using groupby instead of multiple nested loops:

def make_reports(df: DataFrame, report_types: dict):
    cols = ['Date', 'Events', 'CompanyCode']

    for key, val in report_types.items():
        for _, grp in df.query(val['filter']).groupby(cols):
            generate_report(grp, key)
user_128331660
Related