Pandas First Value That Doesnt Return KeyError

36 Views Asked by At

I have a dataframe that has either column A or B

I was hoping for a syntax that uses the or operator to have a stntax like:

Value = df.at['Total','A'] or df.at['Total','B']

but I still receive a KeyError exception. Is there a shorthand way to achieve this instead of doing something like:

if 'A' in df.columns:
    Value = df.at['Total','A']
else:
    Value = df.at['Total','B']
3

There are 3 best solutions below

0
Andrej Kesely On BEST ANSWER

I would use ternary conditional operator:

df.at["Total", "B" if "B" in df else "A"]
1
Sarah On

If your pandas dataframe will always have the same number of columns in the same order, you could use the column index instead of the name. For example, to grab the whole first column:

df[[1]]

Or to grab the column entry in the first row:

df.iloc[0, 1]
0
mozway On

Since you have an inconsistent format, why not make it consistent? rename your columns to have a unique name:

df.rename(columns={'B': 'A'}).at['Total', 'A']

This is handy to generalize to multiple aliases:

aliases = ['B', 'C', 'D']
df.rename(columns=dict.fromkeys(aliases, 'A')).at['Total', 'A']

further generalization

Alternatively, using an index intersection, which also enables you to pass an arbitrary number of aliases:

df.at['Total', next(iter(df.columns.intersection(['A', 'B'])))]

Note that in case of multiple matches (e.g if you have both A and B) this gives you the first match in order of the DataFrame's columns. If you want the first match in a custom order:

# first try to match B, then A
aliases = ['B', 'A']
df.at['Total', next(iter(pd.Index(aliases).intersection(df.columns)))]