Python pandas, .apply, use multiple returned values

Question

in[1]:  

import pandas as pd

    def somefunc(input1, input2):
        output1 = 1
        output2 = 2

        return [output1, output2]


d = {'col1': ['A1', 'B1'], 'col2': ['A2', 'B2']}
df = pd.DataFrame(data=d)

df[['col3', 'col4']] = df.apply(lambda x: somefunc(x['col1'], x['col2']), 
axis=1)

print df  

out[1]:  

   col1 col2  col3  col4
0   A1   A2     1     2
1   B1   B2     1     2

somefunc:
- Gets data from json in url with arguments input1 and input2 and puts it into a new dataframe df2
- Creates output1, output2 with values from df2
- Deletes dataframe when output vaiables is calulated to save RAM

This executes perfectly, but I want to return 3 values which return an error

in[2]:  

import pandas as pd

    def somefunc(input1, input2):
        output1 = 1
        output2 = 2  
        output3 = 3


        return [output1, output2, output3]


d = {'col1': ['A1', 'B1'], 'col2': ['A2', 'B2']}
df = pd.DataFrame(data=d)

df[['col3', 'col4', 'col5']] = df.apply(lambda x: somefunc(x['col1'], x['col2']), 
axis=1)

print df  

out[2]:  
"ValueError: Columns must be same length as key"

The error is different from my actual program, but I guess it correlates;

KeyError: "['col3' 'col4' 'col5'] not in index"

Why does this work with one and two outputs but not three?
Python 2.7.14

jezrael · Accepted Answer · 2018-01-22 14:16:04Z

Solution is return Series from custom function:

def somefunc(input1, input2):
    output1 = 1
    output2 = 2  
    output3 = 3

    return pd.Series([output1, output2, output3])

d = {'col1': ['A1', 'B1'], 'col2': ['A2', 'B2']}
df = pd.DataFrame(data=d)

df[['col3', 'col4', 'col5']] = df.apply(lambda x: somefunc(x['col1'], x['col2']), axis=1)
print (df)
  col1 col2  col3  col4  col5
0   A1   A2     1     2     3
1   B1   B2     1     2     3

Or create DataFrame from lists by constructor:

def somefunc(input1, input2):
    output1 = 1
    output2 = 2  
    output3 = 3

    return [output1, output2, output3]

d = {'col1': ['A1', 'B1'], 'col2': ['A2', 'B2']}
df = pd.DataFrame(data=d)

df1 = df.apply(lambda x: somefunc(x['col1'], x['col2']), axis=1)
df[['col3', 'col4', 'col5']] = pd.DataFrame(df1.values.tolist())
print (df)
  col1 col2  col3  col4  col5
0   A1   A2     1     2     3
1   B1   B2     1     2     3

cs95 · Accepted Answer · 2018-01-22 14:26:16Z

1

You've stumbled upon a shape mismatch between the dataframe and the result.

If you'll note, in the first case, you have a dataframe of 2 columns, and your result is returned as a list of two items. You'll see the intermediate output is this:

   col1  col2
0     1     2
1     1     2

What pandas has done, is that it has special cased the return value, and snapped it into a dataframe of the same size as the input (it assumed that you wanted a dataframe because the shape of the result is the same). Since the result is a DataFrame in this case, you can assign it back.

With the second function, unfortunately, the number of elements returned in the list is 3. This is not the same size as the original dataframe, so what is returned is a single column of lists.

0    [1, 2, 3]
1    [1, 2, 3]
dtype: object

Pandas does not know what to do with this data, and so does not do anything with it. Coincidentally, you cannot assign a Series object of lists to a dataframe slice, and so you get an error.

The solution here, is to not return a list, but instead return a pd.Series object.

def somefunc(input1, input2):
    return pd.Series([1, 2, 3])

This returns a dataframe as a result (pandas now smartly stacks each series to form a row in the final dataframe):

   0  1  2
0  1  2  3
1  1  2  3

And assigning this is as easy as it was before.

edited Jan 22, 2018 at 14:26

answered Jan 22, 2018 at 14:20

cs95

406k106 gold badges744 silver badges797 bronze badges

2 Comments

cs95 Over a year ago

@sinB You're welcome. Although if I were you, I'd be more appreciative of an answer that explains your error PLUS gives you code, instead of just a code based solution. ;)

sinB Over a year ago

You are right @COLDSPEED, newbie mistake. Well, the answer has been chosen this time but I appreciate your effort just as much. next time ;)

Collectives™ on Stack Overflow

Python pandas, .apply, use multiple returned values

2 Answers 2

Comments

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related