I'm trying to write a function that takes the name of a column and splits the dataframe based on the values of that column. I have the following

df_split = df[df.a == 1]

I'm trying to implement the following idea

def f(df,column_name): 
  df_split = df[df.column_name == 1]

Any help is highly appreciated.

CodePudding user response：

Please change the function to following:

def f(df,column_name): 
  df_split = df[df[column_name] == 1]
  return df_split

df.column_name will work only if the dataframe really have a column labelled as column_name so don't use it inside the function