Proper data manipulation script layout, where all merges, drops, aggregations, renames are easily traceable and visible

Viewed 48

I currently have a long script that has one goal: take multiple csv tables of data, merge them into one, while performing various calculations along the way, and then output a final csv table.

I originally had this layout (see LAYOUT A), but found that this made it had to see what columns were being added or merged, because the cleaning and operations methods are listed below everything, so you have to go up and down in the file to see how the table gets altered. This was an attempt to follow the whole keep-things-modular-and-small methodology that I've read about:

import pandas as pd
#...

SOME_MAPPER = {'a':1, 'b':2, ...}
COLUMNS_RENAMER = {'new_col1': 'aaa', ...}

def main():
    df1 = clean_table_1('table1.csv')
    df2 = clean_table_2('table2.csv')
    df3 = clean_table_3('table3.csv')

    df = pd.merge(df1, df2, on='col_a')
    df['new_col1'] = df.apply(lambda r: some_operation(r['x'], r['y'], r['z']), axis=1)
    df['new_col2'] = df['new_col1'].map(SOME_MAPPER)

    df = pd.merge(df, df3, on='new_col2')
    df['new_col3'] = df['something']+df['new_col2']
    df['new_col4'] = df.apply(lambda r: some_other_operation(r['a'], r['b']), axis=1)   

    df = df.rename(columns=COLUMNS_RENAMER)
    return df

def some_operation(x,y,z):
    #<calculations for performing on table column>

def some_other_operation(a,b):
    #<some calculation>

def clean_table_1(fn_1):
    df = pd.read_csv(fn_1)
    df['some_col1'] = 400
    def do_operations_unique_to_table1(df):
        #<operations>
        return df
    df = do_operations_unique_to_table1(df)

    return df

def clean_table_2(fn_2):
    #<similar to clean_table_1>

def clean_table_3(fn_3):
    #<similar to clean_table_1>

if __name__=='__main__':
    main()

My next inclination was to move all the functions in-line with the main script, so its obvious what's being done (see LAYOUT B). This makes it a bit easier to see the linearity of the operations being done, but also makes it a bit messier, so that you can't just quickly read through the main function to get the "overview" of all the operations being done.

# LAYOUT B
import pandas as pd
#...

SOME_MAPPER = {'a':1, 'b':2, ...}
COLUMNS_RENAMER = {'new_col1': 'aaa', ...}

def main():
    def clean_table_1(fn_1):
        df = pd.read_csv(fn_1)
        df['some_col1'] = 400
        def do_operations_unique_to_table1(df):
            #<operations>
            return df
        df = do_operations_unique_to_table1(df)

    df1 = clean_table_1('table1.csv')


    def clean_table_2(fn_2):
        #<similar to clean_table_1>

    df2 = clean_table_2('table2.csv')


    def clean_table_3(fn_3):
        #<similar to clean_table_1>

    df3 = clean_table_3('table3.csv')

    df = pd.merge(df1, df2, on='col_a')

    def some_operation(x,y,z):
        #<calculations for performing on table column>

    df['new_col1'] = df.apply(lambda r: some_operation(r['x'], r['y'], r['z']), axis=1)
    df['new_col2'] = df['new_col1'].map(SOME_MAPPER)

    df = pd.merge(df, df3, on='new_col2')

    def some_other_operation(a,b):
        #<some calculation>

    df['new_col3'] = df['something']+df['new_col2']
    df['new_col4'] = df.apply(lambda r: some_other_operation(r['a'], r['b']), axis=1)   

    df = df.rename(columns=COLUMNS_RENAMER)
    return df

if __name__=='__main__':
    main()

So then I think, well why even have these functions; would it maybe just be easier to follow if it's all at the same level, just as a script like so (LAYOUT C):

# LAYOUT C
import pandas as pd
#...

SOME_MAPPER = {'a':1, 'b':2, ...}
COLUMNS_RENAMER = {'new_col1': 'aaa', ...}

def main():

    df1 = pd.read_csv('table1.csv)
    df1['some_col1'] = 400
    df1 = #<operations on df1>  

    df2 = pd.read_csv('table2.csv)
    df2['some_col2'] = 200
    df2 = #<operations on df2>  

    df3 = pd.read_csv('table3.csv)
    df3['some_col3'] = 800
    df3 = #<operations on df3>  

    df = pd.merge(df1, df2, on='col_a')

    def some_operation(x,y,z):
        #<calculations for performing on table column>

    df['new_col1'] = df.apply(lambda r: some_operation(r['x'], r['y'], r['z']), axis=1)
    df['new_col2'] = df['new_col1'].map(SOME_MAPPER)

    df = pd.merge(df, df3, on='new_col2')

    def some_other_operation(a,b):
        #<some calculation>

    df['new_col3'] = df['something']+df['new_col2']
    df['new_col4'] = df.apply(lambda r: some_other_operation(r['a'], r['b']), axis=1)   

    df = df.rename(columns=COLUMNS_RENAMER)
    return df

if __name__=='__main__':
    main()

The crux of the problem is finding a balance between documenting clearly which columns are being updated, changed, dropped, renamed, merged, etc. while still keeping it modular enough to fit the paradigm of "clean code".

Also, in practice, this script and others are much longer with far more tables being merged into the mix, so this quickly becomes a long list of operations. Should I be breaking up the operations into smaller files and outputting intermediate files or is that just asking to introduce errors? It's a matter of also being able to see all the assumptions made along the way and how they affect the data in its final state, without having to jump between files or scroll way up, way down, etc. to follow the data from A to B, if that makes sense.

If anyone has insights on how to best write these types of data cleaning/manipulation scripts, I would love to hear them.

1 Answers

It is highly subjective topic, but here is my typical approach/remarks/hints:

  • as long as possible optimize for debug/dev time and ease
  • split flow into several scripts (e.g. download, preprocess, ... for every table sepearetely, so for merging every table is prepared separetely)
  • try to keep the same order of operations within script (e.g. type correction, fill na, scaling, new columns, drop columns)
  • for every wrangle script there is load-start and save-end
  • save to pickle (to avoid problems like saving date as string) and small csv (to have easy preview of results outside of python)
  • with "integration points" being data you can easily combine different technologies (caveat: in such case typically you don't use pickle as output but csv/other data format)
  • every script has clearly defined input/output and can be tested and developed separately, also I use asserts on dataframe shapes
  • scripts for visualization/EDA use data from wrangle scripts, but they are never part of wrangle scripts, typically they are also bottleneck
  • combine scripts by e.g. bash if you want simplicity
  • keep length of script below 1 page*
  • *if I have long, convoluted code before I encapsulate it in function, I check if it can be done simpler, in 80% yes but you need to know more about pandas, but you learn something new, pandas doc is typically better, your code is more declarative and idiomatic
  • *if there is no easy way to fix and you use this function in many places put it into utils.py, in docstring put sample >>>f(input) output and some rationale about this function
  • *if function is used across many projects it is worth to make pandas extension like https://github.com/twopirllc/pandas-ta
  • *if I have a lot of columns, I think a lot about hierarchy and groupings and keep it in separate file for every table, for now it is just py file, but I started to consider yaml and a way to document table structure
  • stick to one convention e.g. I don't use inplace=True at all
  • chain operations on dataframe*
  • *if you have a good meaningful name for subchain result that could be used elsewhere, it could be a good place to split script
  • remove main function, if you keep script according to rules above there is nothing wrong in df global variable
  • when I've read from csv, I always check what can be done directly with read_csv parameters e.g. parse date
  • clearly mark 'TEMPORARY HACKS', in long-term they lead to unexpected side-effects
Related