Kristofersen
Kristofersen

Reputation: 2806

Expand a list returned by a function to multiple columns (Pandas)

I have a function that I'm trying to call on each row of a dataframe and I would like it to return 20 different numeric values and each of those be in a separate column of the original dataframe.

For example, this is not the function, but it if this will work the actual one will

def doStuff(x):
    return([x] * 5)

So this will just return the same number 5x. so if I have the dataframe

import pandas as pd

df = pd.DataFrame({'A' : [1,2]})
   A
0  1
1  2
2  3

After calling

df = np.vectorize(doStuff)(df['A'])

It should end up looking like

   A  1  2  3  4  5
0  1  1  1  1  1  1
1  2  2  2  2  2  2
2  3  3  3  3  3  3

Upvotes: 1

Views: 903

Answers (2)

bencekd
bencekd

Reputation: 1605

From pandas 0.23 you can use the result_type argument:

df = pd.DataFrame({'A' : [1,2]})

def doStuff(x):
    return([x] * 5)

df.apply(doStuff, axis=1, result_type='expand')

Upvotes: 1

cs95
cs95

Reputation: 402814

I believe you need df.apply, twice.

In [1254]: df['A'].apply(np.vectorize(doStuff)).apply(pd.Series)
Out[1254]: 
   0  1  2  3  4
0  1  1  1  1  1
1  2  2  2  2  2
2  3  3  3  3  3

You may concatenate this with the original using pd.concat(..., axis=1):

In [1258]: pd.concat([df, df['A'].apply(np.vectorize(doStuff)).apply(pd.Series)], axis=1)
Out[1258]: 
   A  0  1  2  3  4
0  1  1  1  1  1  1
1  2  2  2  2  2  2
2  3  3  3  3  3  3

Upvotes: 2

Related Questions