pandas groupby index: nan conflicts or take only value

Viewed 85

I'm working with some data via pandas which can have some inconsistencies I need to handle. The data is a time series of values:

                     A     B      C
YYYY MM DD hh mm                                
2017 8  20 23 0      1   2.0    NaN
              10     2   4.0    NaN
              20     3   6.0    NaN
              30     4   8.0    NaN
...                ...   ...    ...
2019 6  4  10 10   100   100    NaN
              20   200   102    NaN
              30   300   104    NaN
              40   400   106    NaN
              50   500   108      0

The inconsistency the data can have is multiple entries for the same ["YYYY", "MM", "DD", "hh", "mm"] index. In most cases, the values at that time are are identical so I can use df.drop_duplicates(keep="first") to drop all rows with the same index and column values.

However, there are index collisions where the values are not identical or a non-nan value is only present in one of the rows. The behaviour I'm after is:

For duplicate indices, on a per-column bases:

  1. If only one non-nan value: use that.
  2. If all nan values: use nan.
  3. If all the same (non-nan) value: use that.
  4. If non-equal (non-nan) value: use nan.

For example (is a simplified DataFrame):

     A      B    C
0    1    2.0  NaN
1    2    NaN  NaN
1  100  500.0  NaN
2    3    6.0  NaN
2  200    6.0  NaN
3  300    8.0  NaN
3  300    NaN  5.0
3  300    NaN  NaN
4  400  106.0  NaN

Should result in:

       A      B    C
0    1.0    2.0  NaN
1    NaN  500.0  NaN
2    NaN    6.0  NaN
3  300.0    8.0  5.0
4  400.0  106.0  NaN

I tried to solve this in a couple of methods, but they were both horrendously slow on the dataset size.

current slow solutions (you may need to scroll the code snippet window):

import numpy as np
import pandas as pd

df = pd.DataFrame(
    [
        dict(A=1, B=2.0, C=None),
        dict(A=2, B=None, C=None),
        dict(A=100, B=500, C=None),
        dict(A=3, B=6.0, C=None),
        dict(A=200, B=6.0, C=None),
        dict(A=300, B=8.0, C=None),
        dict(A=300, B=None, C=5.0),
        dict(A=300, B=None, C=None),
        dict(A=400, B=106, C=None),
    ],
    index=[0, 1, 1, 2, 2, 3, 3, 3, 4],
)


# SLOW SOLUTION 1
def canonical(colum_values):
    candidates = colum_values.dropna().unique()
    if len(candidates) == 1:
        return candidates[0]
    else:
        return np.nan


solution_1 = df.groupby(df.index).aggregate(canonical)


# SLOW & UGLY SOLUTION 2
def solve_2(df):
    df = df.copy()

    for dupe in df.index[df.index.duplicated(keep="first")]:
        for column in df.columns:
            values = df[df.index == dupe][column]

            if len(values.dropna().unique()) == 1:
                df.loc[df.index == dupe, column] = values.dropna().iloc[0]
            else:
                df.loc[df.index == dupe, column] = np.nan
    
    # duplicate rows should all now share the same value, so just keep one.
    df.drop_duplicates(keep="first", inplace=True)

    return df


solution_2 = solve_2(df)

Looking for any improvements on these to get better performance on large datasets.

Thanks.

1 Answers

You can group the dataframe on level=0 and aggregate using .first to get the first non-nan value per level=0 group then .mask the values in columns where there are more than one unique values per level=0 group :

g = df.groupby(level=0)
g.first().mask(g.nunique().gt(1))

Details:

g.first()

     A      B    C
0    1    2.0  NaN
1    2  500.0  NaN
2    3    6.0  NaN
3  300    8.0  5.0
4  400  106.0  NaN

g.nunique().gt(1)

       A      B      C
0  False  False  False
1   True  False  False
2   True  False  False
3  False  False  False
4  False  False  False

Result:

       A      B    C
0    1.0    2.0  NaN
1    NaN  500.0  NaN
2    NaN    6.0  NaN
3  300.0    8.0  5.0
4  400.0  106.0  NaN
Related