garej
garej

Reputation: 565

KeyError exception with pandas dataframe colunm

Imaging the one wants to make a function that accepts two kinds of DataFrames - with column "Value" and without. (Say, this is data on assets from different sources and some sources provide no information about value.)

Function foo should run different procedures depending on the dataframe, say, copy "Value" to new column "B" or add new column "B" with 1. A toy example is as follows.

import pandas as pd

df1 = pd.DataFrame({"A":["one", "two", "three"]})
df2 = pd.DataFrame({"A":["one", "two", "three"], "Value":[1, 2, 3]})

What is the best way to handle KeyError: 'Value' in the case of df1.

def foo(df, col_name="Value"):

    <if col_name exists>
        df["B"] = df[col_name] 
    <else>
        df["B"] = 1
    return df

Upvotes: 0

Views: 269

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150785

Try:

df['B'] = df.get('Value', 1)

or equivalently

df['B'] = df['Value'] if 'Value' in df else 1

Upvotes: 3

Related Questions