How to create multiple dataframes using multiple functions

Question

I quite often write a function to return different dataframes based on the parameters I enter. Here's an example dataframe:

np.random.seed(1111)
df = pd.DataFrame({
'Category':np.random.choice( ['Group A','Group B','Group C','Group D'], 10000),
'Sub-Category':np.random.choice( ['X','Y','Z'], 10000),
'Sub-Category-2':np.random.choice( ['G','F','I'], 10000),
'Product':np.random.choice( ['Product 1','Product 2','Product 3'], 10000),
'Units_Sold':np.random.randint(1,100, size=(10000)),
'Dollars_Sold':np.random.randint(100,1000, size=10000),
'Customer':np.random.choice(pd.util.testing.rands_array(10,25,dtype='str'),10000),
'Date':np.random.choice( pd.date_range('1/1/2016','12/31/2018',  
                      freq='M'), 10000)})

I then created a function to perform sub-totals for me like this:

def some_fun(DF1, agg_column, myList=[], *args):
    y = pd.concat([
    DF1.assign(**{x:'[Total]' for x in myList[i:]})\
            .groupby(myList).agg(sumz = (agg_column,'sum')) for i in range(1,len(myList)+1)]).sort_index().unstack(0)
    return y

I then write out lists that I'll pass as arguments to the function:

list_one = [pd.Grouper(key='Date',freq='A'),'Category','Product']
list_two = [pd.Grouper(key='Date',freq='A'),'Category','Sub-Category','Sub-Category-2']
list_three = [pd.Grouper(key='Date',freq='A'),'Sub-Category','Product']

I then have to run each list through my function creating new dataframes:

df1 = some_fun(df,'Units_Sold',list_one)
df2 = some_fun(df,'Dollars_Sold',list_two)
df3 = some_fun(df,'Units_Sold',list_three)

I then use a function to write each of these dataframes to an Excel worksheet. This is just an example - I perform this same exercise 10+ times.

My question - is there a better way to perform this task than to write out df1, df2, df3 with the function information applied? Should I be looking at using a dictionary or some other data type to do this my pythonically with a function?

Answer 1

A dictionary would be my first choice:

variations = ([('Units Sold', list_one), ('Dollars_Sold',list_two), 
              ..., ('Title', some_list)])

df_variations = {}

for i, v in enumerate(variations):
     name = v[0]
     data = v[1]
     df_variations[i] = some_fun(df, name, data)

You might further consider setting the keys to unique / helpful titles for the variations, that goes beyond something like 'Units Sold', which isn't unique in your case.

Answer 2

IIUC,

as Thomas has suggested we can use a dictionary to parse through your data, but with some minor modifications to your function, we can use the dictionary to hold all the required data then pass that through to your function.

the idea is to pass two types of keys, the list of columns and the arguments to your pd.Grouper call.

data_dict = {
    "Units_Sold": {"key": "Date", "freq": "A"},
    "Dollars_Sold": {"key": "Date", "freq": "A"},
    "col_list_1": ["Category", "Product"],
    "col_list_2": ["Category", "Sub-Category", "Sub-Category-2"],
    "col_list_3": ["Sub-Category", "Product"],
}

---

def some_fun(dataframe, agg_col, dictionary,column_list, *args):

    key = dictionary[agg_col]["key"]

    frequency = dictionary[agg_col]["freq"]

    myList = [pd.Grouper(key=key, freq=frequency), *dictionary[column_list]]

    y = (
        pd.concat(
            [
                dataframe.assign(**{x: "[Total]" for x in myList[i:]})
                .groupby(myList)
                .agg(sumz=(agg_col, "sum"))
                for i in range(1, len(myList) + 1)
            ]
        )
        .sort_index()
        .unstack(0)
    )
    return y

---

Test.

df1 = some_fun(df,'Units_Sold',data_dict,'col_list_3')
print(df1)
                                 sumz                      
Date                   2016-12-31 2017-12-31 2018-12-31
Sub-Category Product                                   
X            Product 1      18308      17839      18776
             Product 2      18067      19309      18077
             Product 3      17943      19121      17675
             [Total]        54318      56269      54528
Y            Product 1      20699      18593      18103
             Product 2      18642      19712      17122
             Product 3      17701      19263      20123
             [Total]        57042      57568      55348
Z            Product 1      19077      17401      19138
             Product 2      17207      21434      18817
             Product 3      18405      17300      17462
             [Total]        54689      56135      55417
[Total]      [Total]       166049     169972     165293

as you want to automate the writing of the 10x worksheets, we can again do that with a dictionary call over your function:

matches = {'Units_Sold': ['col_list_1','col_list_3'],
          'Dollars_Sold' : ['col_list_2']}

then a simple for loop to write all the files to a single excel sheet, change this to match your required behavior.

writer = pd.ExcelWriter('finished_excel_file.xlsx')
for key,value in matches.items():
    for items in value:        
        dataframe = some_fun(df,k,data_dict,items)
        dataframe.to_excel(writer,f'{key}_{items}')
writer.save()