Using pandas apply with if condition

Question

below I've tried to give an example of my problem. I not sure I understand the following error I receive right. f['C'], df['D'] = zip(*df.apply(lambda x: 0 if x['A'] == 1 else some_func(x['A'], x['B']), axis=1)) TypeError: 'int' object is not iterable

Is it because if in this case the 0 which is return when x['A'] == 1 is only a single int and can, therefore, not be unzipped? And of course any help solving this error would be much appreciated. Cheers.

import pandas as pd

# A function that returns multiple things.
def some_func(x,y):
    return x+y, x-y

# Example DataFrame
df = pd.DataFrame({'A': range(5), 'B':range(5,10,1)})

# Example usage.
df['C'], df['D'] = zip(*df.apply(lambda x: 0 if x['A'] == 1 else some_func(x['A'], x['B']), axis=1))

print(df)

Answer 1

You can do two things.

One, as Parfait suggested in the comments, you return a tuple instead of an integer - lambda x: (0, 0) if x['A'] == 1 else some_func(x['A'], x['B']) .

Else, if your logic is more complex, I suggest you move your if logic to the function itself.

# A function that returns multiple things.
def some_func(df):
    if df['A'] == 1: # Replace with your condition
        return (None, None) # Return a tuple, with whatever values. I suggest None, but 0 is OK.
    # Some elif logic
    else:
         return x+y, x-y

Then:

df['C'], df['D'] = zip(*df.apply(some_func, axis=1))