Pandas: Loop over GroupBy / Two aggregations / Fillna not inplace - pandas

Question

I have a dataframe which uses an identifier for groups and has several columns with missing values.
toy_df = pd.DataFrame({'ID':[1,1, 1, 2, 2, 2, 3, 3, 3],
'Age': [10, 10, 10, 20, 20, 20, 20, 20, 20],
'A': [np.nan, 5, 5, np.nan, np.nan, np.nan, 10, 12, 12],
'B': [3, 4, 5, 2, 2, 1, np.nan, 4, 3]})
ID Age A B
0 1 10 NaN 3.0
1 1 10 5.0 4.0
2 1 10 5.0 5.0
3 2 20 NaN 2.0
4 2 20 NaN 2.0
5 2 20 NaN 1.0
6 3 20 10.0 NaN
7 3 20 12.0 4.0
8 3 20 12.0 3.0
Now I want to fill NaNs by some rules either within groups of same age or just within the ID group
group_mode = toy_df.groupby('Age')['A'].apply(lambda x: list(x.mode()))
group_median = toy_df.groupby('Age')['A'].median()
def impute_column(series, group_mode, group_median, agg_key, key, age):
if series.isna().sum() == series.shape[0]:
modes = group_mode[group_mode.index == age]
#if multiple modes are available use median
if np.ravel(modes.to_list()).shape[0] > 1:
median_ = group_median[group_median.index == age]
series = series.fillna(value = median_)
else:
mode_ = modes.item()[0]
series = series.fillna(value = mode_)
#if up to 3 values are missing use linear interpolation
elif series.isna().sum() < 4:
series = series.interpolate(limit_direction='both', method='linear')
#else we have sparse values / use median
else:
median_ = series.median()
series = series.fillna(value = median_)
return series
And if I test it with one of the columns & groups it works
impute_column(series=toy_df['A'], group_mode=group_mode, group_median=group_median,
agg_key='Age', key = 'A', age=10)
0 10.0
1 5.0
2 5.0
3 10.0
4 10.0
5 10.0
6 10.0
7 12.0
8 12.0
Now I want to be efficient and
group over IDs
loop over the grouped object
loop over all columns
and update my dataframe
all_columns = ['A', 'B']
grouped = toy_df.groupby('ID')
for key in all_columns:
group_mode = toy_df.groupby('Age')[key].apply(lambda x: list(x.mode()))
group_median = toy_df.groupby('Age')[key].median()
for _, group in grouped:
age = group['Age'].iloc[0]
group[key] = impute_column(series=group[key], group_mode=group_mode, group_median=group_median,
agg_key='Age', key=key, age=age)
The calculations are running (I've printed them out), but the final dataframe ain't updated
ID Age A B
0 1 10 NaN 3.0
1 1 10 5.0 4.0
2 1 10 5.0 5.0
3 2 20 NaN 2.0
4 2 20 NaN 2.0
5 2 20 NaN 1.0
6 3 20 10.0 NaN
7 3 20 12.0 4.0
8 3 20 12.0 3.0
what seems to work is the code below.
but as you can see, it does not follow the bullet points above. Further I am pretty sure, that computing big groupby objects for each group is immensely inefficient
def impute_column(series, group_mode, group_median, key, age):
if series.isna().sum() == series.shape[0]:
modes = group_mode[group_mode.index == age]
#if multiple modes are available use median
if np.ravel(modes.to_list()).shape[0] > 1:
median_ = group_median[group_median.index == age]
series = series.fillna(value = median_)
else:
mode_ = modes.item()[0]
series = series.fillna(value = mode_)
#if up to 3 values are missing use linear interpolation
elif series.isna().sum() < 4:
series = series.interpolate(limit_direction='both', method='linear')
#else we have sparse values / use median
else:
median_ = series.median()
series = series.fillna(value = median_)
return series
def impute_frame(data, full_data, agg_key):
age = data['Age'].iloc[0]
for key in ['A', 'B']:
group_mode = full_data.groupby(agg_key)[key].apply(lambda x: list(x.mode()))
group_median = full_data.groupby(agg_key)[key].median()
data[key] = impute_column(data[key], group_mode, group_median, key, age)
return data
toy_df.groupby('ID').apply(impute_frame, full_data=toy_df, agg_key='Age')
ID Age A B
0 1 10 5.0 3.0
1 1 10 5.0 4.0
2 1 10 5.0 5.0
3 2 20 12.0 2.0
4 2 20 12.0 2.0
5 2 20 12.0 1.0
6 3 20 10.0 4.0
7 3 20 12.0 4.0
8 3 20 12.0 3.0