Replace null values with mean for specific columns in pyspark - replace

Question

I would like to replace null values with mean for the age and height column. I know there is a post
Fill Pyspark dataframe column null values with average value from same column
but in this post the given function throws an error.
df = spark.createDataFrame([(1, 'John', 1.79, 28,'M', 'Doctor'),
(2, 'Steve', 1.78, 45,'M', None),
(3, 'Emma', 1.75, None, None, None),
(4, 'Ashley',1.6, 33,'F', 'Analyst'),
(5, 'Olivia', 1.8, 54,'F', 'Teacher'),
(6, 'Hannah', 1.82, None, 'F', None),
(7, 'William', 1.7, 42,'M', 'Engineer'),
(None,None,None,None,None,None),
(8,'Ethan',1.55,38,'M','Doctor'),
(9,'Hannah',1.65,None,'F','Doctor')]
, ['Id', 'Name', 'Height', 'Age', 'Gender', 'Profession'])
the function in the post given
def fill_with_mean(df, exclude=set()):
stats = df.agg(*(
avg(c).alias(c) for c in df.columns if c not in exclude
))
return df.na.fill(stats.first().asDict())
fill_with_mean(df, ["Age", "Height"])
when I run this function, it says
NameError: name 'avg' is not defined
Can anybody fix this? Thank you.

Answer 1

Fixed example. It works for me in a way as you expect!
from pyspark.sql.functions import avg
df = spark.createDataFrame(
[
(1, 'John', 1.79, 28, 'M', 'Doctor'),
(2, 'Steve', 1.78, 45, 'M', None),
(3, 'Emma', 1.75, None, None, None),
(4, 'Ashley', 1.6, 33, 'F', 'Analyst'),
(5, 'Olivia', 1.8, 54, 'F', 'Teacher'),
(6, 'Hannah', 1.82, None, 'F', None),
(7, 'William', 1.7, 42, 'M', 'Engineer'),
(None, None, None, None, None, None),
(8, 'Ethan', 1.55, 38, 'M', 'Doctor'),
(9, 'Hannah', 1.65, None, 'F', 'Doctor')
],
['Id', 'Name', 'Height', 'Age', 'Gender', 'Profession']
)
def fill_with_mean(this_df, exclude=set()):
stats = this_df.agg(*(avg(c).alias(c) for c in this_df.columns if c not in exclude))
return this_df.na.fill(stats.first().asDict())
res = fill_with_mean(df, ["Gender", "Profession", "Id", "Name"])
res.show()