Search code examples
pythonpandasdataframeoperators

Pass operator to function that selects on dataframe


How to pass operator (eq, ne, gt etc.) to dataframe select function?

import pandas as pd
from operator import eq, ne

df = pd.DataFrame({'Greet Ing': ['Hi Friend', 123, 'Hi Friend', '456']})

df[df['Greet Ing'] == 'Hi Friend']
df.query("`Greet Ing` == 'Hi Friend'")


def filter(dataframe, column, operator, value):
    return dataframe[dataframe[column].operator(value)]

filter(df, 'Greet Ing', eq, 'Hi Friend')  # Not working

The lazy version would be treat operator as string and conditional.

if o == 'eq': ..
elif o == 'ne': ..

Solution

  • Create function with pass Series to operator:

    from operator import eq, ne
    
    def f(df1, c, op, y):
        return df1[op(df1[c], y)]
    
    df2 = f(df, 'Greet Ing', eq, 'Hi Friend')
    print (df2)
       Greet Ing
    0  Hi Friend
    2  Hi Friend
    

    Or you can map operator by dictionary:

    from operator import eq, ne
    
    ops = {'==': eq, '!=': ne}
    
    def f(df1, c, op, y):
        return df1[ops[op](df1[c], y)]
    
    df2 = f(df, 'Greet Ing', '==', 'Hi Friend')
    print (df2)
       Greet Ing
    0  Hi Friend
    2  Hi Friend