Creating multiple pandas columns from function returning dict

Question

I have a function that returns a dict object and I would like to take advantage of pandas/numpy's ability to perform columnwise operations/vectorization for this function across every row of a dataframe. The inputs for the function are specified in the dataframe and I want the outputs of the function to become new columns on the existing dataframe. Below is an example.

def func(a, b, c):
    return {
        "a_calc": a * 2, 
        "b_calc": b * 3, 
        "c_calc": c * 4
    }

df = pd.DataFrame([{"a":1, "b":2, "c": 3}, {"a": 4, "b": 5, "c": 6}])
   a  b  c
0  1  2  3
1  4  5  6

Desired Output:

   a  b  c  a_calc  b_calc  c_calc
0  1  2  3       2       6      12
1  4  5  6       8      15      24

I was reading this answer and it got most of the way there but I couldn't quite figure out how to do it for when the function returns a dict object with the desired column names as the keys within the dict.

Scott Boston · Accepted Answer · 2021-04-07 02:06:29Z

4

Let's use some dataframe unpacking:

df.join(pd.DataFrame(func(**df)))

Output:

   a  b  c  a_calc  b_calc  c_calc
0  1  2  3       2       6      12
1  4  5  6       8      15      24

Or really getting cute:

df.assign(**func(**df))

answered Apr 7, 2021 at 2:06

Scott Boston

154k15 gold badges160 silver badges207 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Quang Hoang · Accepted Answer · 2021-04-07 02:20:18Z

2

If you cannot modify your function, you can do:

df.join(pd.DataFrame(func(df['a'], df['b'],df['c']), index=df.index))

Output:

   a  b  c  a_calc  b_calc  c_calc
0  1  2  3       2       6      12
1  4  5  6       8      15      24

Note We exploit the fact that func can accept series input and works in parallel. In the general case, you need a for loop:

pd.DataFrame([func(x['a'], x['b'], x['c']) for _, x in df.iterrows()],
              index=df.index)

edited Apr 7, 2021 at 2:20

answered Apr 7, 2021 at 1:56

Quang Hoang

151k11 gold badges64 silver badges86 bronze badges

6 Comments

Austin Ulfers Over a year ago

In my case, func() cannot accept series data. Is there any other way besides using .iterrows()?

Scott Boston Over a year ago

@AustinUlfers Can you rewrite your function to accept a pd.Series?

Austin Ulfers Over a year ago

@ScottBoston Its a relatively long function so I would like to avoid having to rewrite it if possible.

Scott Boston Over a year ago

@AustinUlfers I think the only other way is to use lambda with axis=1 which is woefully slow and inefficient.

Quang Hoang Over a year ago

@AustinUlfers iterrows, apply, vanilla for loop are essentially equivalent. You could avoid one or another, but in general you would need to use one.

|

Collectives™ on Stack Overflow

Creating multiple pandas columns from function returning dict

2 Answers 2

Comments

6 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

6 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related