Add columns to Dataframe when apply custom function that returns dictionary

Question

def tFunc(row):
    if (random.random()>0.5):
        info={'A': 'a', 'B': 'b', 'C': 'c', 'D': 'd'}
    else:
        info={'A': 'a', 'B': 'b', 'C': 'c'}
    return info
    # Workaround
    # for key in info:
    #     row[key]=info[key]
    # return row

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

Consider an existing df. tFunc is a function that returns a dictionary info. Want to expand the existing df. Is there any better/faster way than the workaround method?

Answer 1

You can join the output of your apply call:

out = df.join(df.apply(tFunc, axis=1, result_type='expand'))

example with an input DataFrame that has one column "col":

   col  A  B  C    D
0    0  a  b  c  NaN
1    1  a  b  c    d
2    2  a  b  c    d

Another (not recommended) option might be to combine_first the dictionary as Series to the input row and return this. The drawback is that it might affect the original dtypes:

def tFunc(row):
    if (random.random()>0.5):
        info={'A': 'a', 'B': 'b', 'C': 'c', 'D': 'd'}
    else:
        info={'A': 'a', 'B': 'b', 'C': 'c'}
    return row.combine_first(pd.Series(info))

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

example output:

   A  B  C    D  col
0  a  b  c  NaN  0.0
1  a  b  c  NaN  1.0
2  a  b  c    d  2.0