conditional_col_creation_benchmark.py

import pandas as pd
import perfplot

def numpy_where(df):
  return df.assign(is_rich=np.where(df['salary'] >= 50, 'yes', 'no'))
  
def list_comp(df):
  return df.assign(is_rich=['yes' if x >= 50 else 'no' for x in df['salary']])

def loc(df):
  df = df.assign(is_rich='no')
  df.loc[df['salary'] > 50, 'is_rich'] = 'yes'
  return df

kernels = [numpy_where, list_comp, loc]

df = pd.DataFrame({
  'age' : [21, 45, 45, 5],
  'salary' : [20, 40, 10, 100]
})

perfplot.show(
    setup=lambda n: pd.concat([df] * n, ignore_index=True),
    kernels=kernels,
    labels=[f.__name__ for f in kernels],
    n_range=[2**k for k in range(0, 20)],
    logx=True,
    logy=True,
    xlabel='N',
    equality_check=pd.DataFrame.equals)