[英] (py)spark weighted average taking account of missing values
有没有一种标准的方法来计算pyspark中的加权平均值,忽略分母和中的缺失值?
以下面的例子为例:
# create data
data2 = [(1,1,1,1),
(1,None,1,2),
(2,1,1,1),
(2,3,1,2),
]
schema = (StructType([
StructField("group",IntegerType(),True),
StructField("var1",IntegerType(),True),
StructField("var2",IntegerType(),True),
StructField("wght", IntegerType(), True),
]))
df = spark.createDataFrame(data=data2,schema=schema)
df.printSchema()
df.show(truncate=False)
+-----+----+----+----+
|group|var1|var2|wght|
+-----+----+----+----+
|1 |1 |1 |1 |
|1 |null|1 |2 |
|2 |1 |1 |1 |
|2 |3 |1 |2 |
+-----+----+----+----+
我可以计算其他地方记录的加权平均值:
(df.groupBy("group").agg(
(F.sum(col("var1")*col("wght"))/F.sum("wght")).alias("wgtd_var1"),
(F.sum(col("var2")*col("wght"))/F.sum("wght")).alias("wgtd_var2")).show(truncate=False))
+-----+------------------+---------+
|group|wgtd_var1 |wgtd_var2|
+-----+------------------+---------+
|1 |0.3333333333333333|1.0 |
|2 |2.3333333333333335|1.0 |
+-----+------------------+---------+
但问题是,对于第一组,Weighted average应该是1,因为第二次观察不能使用.我可以
# get new weights
df = (df.withColumn("wghtvar1", F.when(col("var1").isNull(), None)
.otherwise(col("wght")))
.withColumn("wghtvar2", F.when(col("var2").isNull(), None)
.otherwise(col("wght"))))
# compute correct weighted average
(df.groupBy("group").agg(
(F.sum(col("var1")*col("wghtvar1"))/F.sum("wghtvar1")).alias("wgtd_var1"),
(F.sum(col("var2")*col("wghtvar2"))/F.sum("wghtvar2")).alias("wgtd_var2")).show(truncate=False))
+-----+------------------+---------+
|group|wgtd_var1 |wgtd_var2|
+-----+------------------+---------+
|1 |1.0 |1.0 |
|2 |2.3333333333333335|1.0 |
+-----+------------------+---------+
有没有一种规范的方法可以做到这一点?