iterating through pandas dataframe and create new columns using custom function

Question

I have a pandas dataframe which (obviously) contain some data. I have created a function that outputs a number new columns. How can I iterate or apply that function?

I have created a minimum example below ( not the actual problem), with a dataframe and function.

EDIT: Think of the function as a "black box". We don't now what is in, but based on the input it returns a dataframe, that should be added to the existing dataframe.

import pandas as pd
a=pd.DataFrame({"input1": ["a","b"], "input2":[3,2]})

  input1  input2
0      a       3
1      b       2

def f(i1,i2):
    return(pd.DataFrame([{"repeat" : [i1]*i2, "square":i2**2 }]))

So in this case the function returns two new columns "repeat" and "square"

f(a.iloc[0,0],a.iloc[0,1])

      repeat  square
0  [a, a, a]       9

f(a.iloc[1,0],a.iloc[1,1])
   repeat  square
0  [b, b]       4

What I would like to end up with a data frame like this

  input1  input2     repeat  square
0      a       3  [a, a, a]       9
1      b       2     [b, b]       4

Does anyone have an elegant solution to this?

Answer 1

I'd do it using assign :

a = a.assign(
    repeat = a['input1'].repeat(a['input2']).groupby(level=0).agg(list),
    square = np.square(a['input2']),
)

Output:

>>> a
  input1  input2     repeat  square
0      a       3  [a, a, a]       9
1      b       2     [b, b]       4

Answer 2

You can try this modification of the f function:

import pandas as pd


def f(df, col1, col2):
    df_ = pd.DataFrame([{"repeat": list(df[col1] * df[col2])}]).explode("repeat", ignore_index=True)
    df_["square"] = list(df[col2] ** 2)
    return pd.concat([df, df_], axis=1)


a = pd.DataFrame({"input1": ["a", "b"], "input2": [3, 2]})
f(a, "input1", "input2")

Answer 3

How about using pd.concat ?

generated_df = pd.concat([f(*args) for args in a.to_numpy()], ignore_index=True)
out = pd.concat([a, generated_df], axis=1)

Output:

>>> out
  input1  input2     repeat  square
0      a       3  [a, a, a]       9
1      b       2     [b, b]       4