Creating a function to filter a dataframe based on a dictionary

Viewed 21

I am new coding in Python and I'm facing an issue.

I've a pandas dataframe like this:

df = pd.DataFrame({'color':['red', 'green', 'blue'], 'brand':['Ford','fiat', 'opel'], 'year':[2016,2016,2017]})

Which I want to filter by a dictionary like this:

d = {'color':'red', 'year':2016}

I can filter using this:

df.loc[np.all(df[list(d)] == pd.Series(d), axis=1)]

But when I try to put this last statement on a function, it doesn't work anymore and I get "none" as result. Anyone knows where are my mistakes?

my code:

enter image description here

1 Answers

There are 2 ways to get the output:

(1) print it in the function. Use this if you only need to visualise or to check the output:

def filtro(df, **d):
    df = df.loc[np.all(df[list(d)] == pd.Series(d), axis=1)]
    print(df)

filtro(df, **d)

  color brand  year
0   red  Ford  2016

(2) use return statement (as suggested by @mozway and @NuriTaş. Use this if you need the output for subsequent steps or downstream operations:

def filtro(df, **d):
    df = df.loc[np.all(df[list(d)] == pd.Series(d), axis=1)]
    return df

df1 = filtro(df, **d)
print(df1)

  color brand  year
0   red  Ford  2016
Related