заменяя NA значением из списка Pyspark Dataframe

Question

У меня есть искровой датафрейм с 20 столбцами. Я хотел бы заменить значения NA в выбранных столбцах (числовые столбцы) их средними.

У меня есть список имен числовых столбцов и список их средних. Я написал нижеприведенную функцию, но я не уверен, как применить ее к Dataframe

NumColNames=['MinTemp','MaxTemp','Rainfall','WindGustSpeed',\
             'WindSpeed9am', 'WindSpeed3pm', 'Humidity9am', 'Humidity3pm', 'Pressure9am','Pressure3pm']

avgLst=[12,33,44,21,132,35,22,11,4,55]

def replaceNaNum(df, NumColNames,avgLst):
    #iterate through numerical columns names
    for column in NumColNames:
        #iterate through the aveagres in avgLst
        for avg in avgLst:
            #replace each NA value in every column with the corresponding average 
            df=df.withColumn(column, when(df[column] == 'NA',\
                                                       avg).otherwise(df[column]))
    return df

Любой вклад приветствуется, спасибо

Answer 1

Вы можете использовать zip здесь, чтобы связать имена столбцов и соответствующие средние, а затем извлечь их из одного цикла:

for column, avg in zip(NumColNames, avgLst):
    df = df.withColumn(column, when(df[column] == 'NA',\
                                                   avg).otherwise(df[column]))