pandas apply function (UDF) fails to return multiple values in case the dataframe is empty

user_2587904 0
Viewed 117

I want to be able to return more than a single columns from a pandas UDF (apply function). This works nicely, as long as the data frame is non empty! In case it is empty it fails with: not enough values to unpack (expected 3, got 0). Is this to be considered as a bug in pandas? Or should the user be forced to manually check the length of the filtered data frame before executing the function? Or is there a better way not to run into this problem?

import pandas as pd

df = pd.DataFrame({'foo':[1,2,3], 'bar':[4,5,6]})

def my_function(x):
    #print(x)
    # some computation
    # returns multiple values (tuple)
    # simplified here
    return 1,1,1

df = df[df.foo > 10]
df['r1'], df['r2'], df['r3'] = zip(*df.apply(my_function, axis=1))
df
pandas apply user-defined-functions python
1 Answers

One solution is to use pd.concat in combination with result_type='expand'.

cols = {0: 'r1', 1: 'r2', 2: 'r3'}
df = pd.concat([df, df.apply(my_function, axis=1, result_type='expand')], axis=1).rename(columns=cols)

You have to rename the columns afterwards. Also, the resulting empty dataframe repeats the first two columns:

Output:

    foo bar foo bar

versus

    foo bar

Both dataframes are empty, so it may be of no interest for you.

I think to check for empty dataframes in pandas is good practice. So, Siddhants solution in the comment is just fine.

user_85733360
Related