Search code examples
pythonpandasdataframeapply

apply list-returning function to all rows in a pandas DataFrame


I need to apply a function and have the returned list's values inserted into columns in the (new) dataframe. That is, if I have this:

import pandas as pd
def fu(myArg):
    return {"one":myArg,"two":20}
li = [10,20]
df = pd.DataFrame(li,columns=["myArg"])

i would like to apply the fu function to each row of the dataframe so that the output is :

    myArg   one two
0   10      10  20
1   20      20  20

(sorry for the extra spaced there, added for display reasons).

How do I go about doing this, efficiently?


Solution

  • Get your function to return a Series, apply, then join:

    def fu(myArg):
        return pd.Series({'one': myArg, 'two': 20})
    
    out = df.join(df['myArg'].apply(fu))
    

    If you can't modify the function:

    def fu(myArg):
        return {'one': myArg, 'two': 20}
    
    out = df.join(df['myArg'].apply(lambda x: pd.Series(fu(x))))
    

    Or convert the Series of dictionaries to DataFrame with json_normalize (warning, json_normalize doesn't maintain the index, you need to set it manually!):

    def fu(myArg):
        return {'one': myArg, 'two': 20}
    
    out = df.join(pd.json_normalize(df['myArg'].apply(fu)).set_index(df.index))
    

    Output:

       myArg  one  two
    0     10   10   20
    1     20   20   20