Использование PySpark Imputer для сгруппированных данных

Question

У меня есть столбец Class, который может быть 1, 2 или 3, и другой столбец Age с некоторыми пропущенными данными. Я хочу вменять в среднем Age каждой Class группы.

Я хочу сделать что-то вместе:

grouped_data = df.groupBy('Class')
imputer = Imputer(inputCols=['Age'], outputCols=['imputed_Age'])
imputer.fit(grouped_data)

Есть ли обходной путь к этому?

Спасибо за ваше время

Answer 1

Используя Imputer, вы можете отфильтровать набор данных до каждого значения Class, вычислить среднее значение, а затем соединить их обратно, так как вы заранее знаете, какими могут быть значения:

subsets = []
for i in range(1, 4):
    imputer = Imputer(inputCols=['Age'], outputCols=['imputed_Age'])
    subset_df = df.filter(col('Class') == i)
    imputed_subset = imputer.fit(subset_df).transform(subset_df)
    subsets.append(imputed_subset)
# Union them together
# If you only have 3 just do it without a loop
imputed_df = subsets[0].unionByName(subsets[1]).unionByName(subsets[2])

Если вы заранее не знаете, что это за значения, или если их нелегко повторить, вы можете groupBy, получить средние значения для каждой группы в виде DataFrame, а затем объединить их обратно в исходный dataframe.

import pyspark.sql.functions as F
averages = df.groupBy("Class").agg(F.avg("Age").alias("avgAge"))
df_with_avgs = df.join(averages, on="Class")
imputed_df = df_with_avgs.withColumn("imputedAge", F.coalesce("Age", "avgAge"))

Answer 2

Вам необходимо преобразовать ваш фрейм данных с помощью подходящей модели. Затем возьмите среднее значение по заполненным данным:

from pyspark.sql import functions as F

imputer = Imputer(inputCols=['Age'], outputCols=['imputed_Age'])
imp_model = imputer.fit(df)
transformed_df = imp_model.transform(df)
transformed_df \
    .groupBy('Class') \
    .agg(F.avg('Age'))