Pandas: imputing descriptive stats using a groupby with a variable - pandas

Question

I have a data frame like this:
input_df = pd.DataFrame({"sex": ["M", "F", "F", "M", "M"], "Class": [1, 2, 2, 1, 1], "Age":[40, 30, 30, 50, NaN]})
What I want to do is to impute the missing value for the age based on the sex and class columns.
I have tried doing it with a function, conditional_impute. What the function does is take a data frame and a condition and then use it to impute the age based on the sex and class grouping. Butthe caveat is that the condition can either be a mean or median and if not either of these two, the function has to raise an error.
So I did this:
### START FUNCTION
def conditional_impute(input_df, choice='median'):
my_df = input_df.copy()
# if choice is not median or mean, raise valueerror
if choice == "mean" or choice == "median":
my_df['Age'] = my_df['Age'].fillna(my_df.groupby(["Sex","Pclass"])['Age'].transform(choice))
else:
raise ValueError()
# round the values in Age colum
my_df['Age'] = round(my_df['Age'], 1)
return my_df
### END FUNCTION
But I am getting an error when I call it.
conditional_impute(train_df, choice='mean')
What could I possibly be doing wrong? I really cannot get a handle on this.

Answer 1

If you give the right inputs, it outputs just fine...
# Fixed input to match function:
df = pd.DataFrame({"Sex": ["M", "F", "F", "M", "M"], "Pclass": [1, 2, 2, 1, 1], "Age":[40, 30, 30, 50, np.nan]})
def conditional_impute(input_df, choice='median'):
my_df = input_df.copy()
# if choice is not median or mean, raise valueerror
if choice == "mean" or choice == "median":
my_df['Age'] = my_df['Age'].fillna(my_df.groupby(["Sex","Pclass"])['Age'].transform(choice))
else:
raise ValueError()
# round the values in Age colum
my_df['Age'] = round(my_df['Age'], 1)
return my_df
conditional_impute(df, choice='mean')
Output:
Sex Pclass Age
0 M 1 40.0
1 F 2 30.0
2 F 2 30.0
3 M 1 50.0
4 M 1 45.0