Replace null values with mean for specific columns in pyspark

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.

did you import this function? `from pyspark.sql.functions import avg` — vvg, Aug 02 '18 at 15:13
No below is the error --------------------------------------------------------------------------- Py4JJavaError Traceback (most recent call last) in () ----> 1 fill_with_mean(df, ["Age", "Height"]) in fill_with_mean(df, exclude) 4 avg(c).alias(c) for c in df.columns if c not in exclude 5 )) ----> 6 return df.na.fill(stats.first().asDict()) — melik, Aug 02 '18 at 15:35

score 1 · Accepted Answer · answered Aug 02 '18 at 15:36

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()

Replace null values with mean for specific columns in pyspark

NameError: name 'avg' is not defined

1 Answers1