[SOLVED] iterating through pandas dataframe and create new columns using custom function


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"


      repeat  square
0  [a, a, a]       9

   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?


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)


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

Answered By – richardec

Answer Checked By – Terry (BugsFixing Volunteer)

Leave a Reply

Your email address will not be published. Required fields are marked *