Pyspark Dataframe Imputations - 替换未知&缺少值与列平均值基于指定条件

时间:2016-05-24 22:29:10

标签: python replace pyspark aggregation imputation

给定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方法替换未知/缺失值。

我尝试过使用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循环'中并循环遍历所有适用的列来执行此操作吗?

更新

还想出了如何遍历列...这是一个例子。

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

1 个答案:

答案 0 :(得分:7)

如果未知或缺失的年龄是某种值:

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"])