Pyspark Dataframe 插补——根据指定条件用列平均值替换未知值和缺失值

Question

给定一个 Spark 数据框，我想根据该列的非缺失值和非未知值计算列均值。然后我想采用这个意思并用它来替换列的缺失值和未知值。

例如，假设我正在使用：

名为 df 的数据框，其中每条记录代表一个人，所有列都是整数或数字
名为年龄的列（每条记录的年龄）
名为 missing_age 的列（如果该人没有年龄则等于 1，否则为 0）
名为 unknown_age 的列（如果该人的年龄未知，则等于 1，否则为 0）

然后我可以计算这个平均值，如下所示。

calc_mean = df.where((col("unknown_age") == 0) & (col("missing_age") == 0))
.agg(avg(col("age")))

或通过 SQL 和 windows 函数，

mean_compute = hiveContext.sql("select avg(age) over() as mean from df 
where missing_age = 0 and unknown_age = 0")

如果可以的话，我不想使用SQL/windows功能。我的挑战是采用这种方法并使用非 SQL 方法用它替换 unknown/missing 值。

我尝试过使用 when()、where()、replace()、withColumn、UDF 和组合...无论我做什么，我要么会出错，要么结果不是我所期望的.这是我尝试过但没有用的许多事情之一的例子。

imputed = df.when((col("unknown_age") == 1) | (col("missing_age") == 1),
calc_mean).otherwise("age")

我在网上搜索过，但没有找到类似的插补类型问题，因此非常感谢您的帮助。这可能是我错过的非常简单的事情。

旁注 -- 我正在尝试将此代码应用于 Spark Dataframe 中列名称中没有 unknown_ 或 missing_ 的所有列。我可以将 Spark 相关代码包装在 Python 'for loop' 中并遍历所有适用的列来执行此操作吗？

更新：

还弄清楚了如何遍历列...这是一个示例。

for x in df.columns:
    if 'unknown_' not in x and 'missing_' not in x:
        avg_compute = df.where(df['missing_' + x] != 1).agg(avg(x)).first()[0]
        df = df.withColumn(x + 'mean_miss_imp', when((df['missing_' + x] == 1), 
        avg_compute).otherwise(df[x]))

Answer 1

如果未知或缺失的年龄是某个值：

from pyspark.sql.functions import col, avg, when

df = sc.parallelize([
    (10, 0, 0), (20, 0, 0), (-1, 1, 0), (-1, 0, 1)
]).toDF(["age", "missing_age", "unknown_age"])

avg_age = df.where(
    (col("unknown_age") != 1) & (col("missing_age") != 1)
).agg(avg("age")).first()[0]

df.withColumn("age_imp", when(
    (col("unknown_age") == 1) | (col("missing_age") == 1), avg_age
).otherwise(col("age")))

如果未知或失踪的年龄为 NULL，您可以将其简化为：

df = sc.parallelize([
    (10, 0, 0), (20, 0, 0), (None, 1, 0), (None, 0, 1)
]).toDF(["age", "missing_age", "unknown_age"])

df.na.fill(df.na.drop().agg(avg("age")).first()[0], ["age"])

Pyspark Dataframe 插补——根据指定条件用列平均值替换未知值和缺失值

Pyspark Dataframe Imputations -- Replace Unknown & Missing Values with Column Mean based on specified condition

python

replace

aggregation

pyspark

imputation