是否有一種規范的方法來計算 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 |
----- ------------------ ---------
但問題是對于第 1 組，加權平均值應該是 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 |
----- ------------------ ---------
有沒有規范的方法來做到這一點？
uj5u.com熱心網友回復：
差別不大，但至少這可以使您免于為每個變數創建新的 wght 列。
條件聚合。
df = (df.groupby('group')
.agg(
(F.sum(F.when(F.col('var1').isNotNull(), F.col('var1') * F.col('wght')))
/
(F.sum(F.when(F.col('var1').isNotNull(), F.col('wght'))))
).alias('wgtd_var1')
))
要將其應用于多個var，我可以使用串列理解。
df = (df.groupby('group')
.agg(*[
(F.sum(F.when(F.col(x).isNotNull(), F.col(x) * F.col('wght')))
/
(F.sum(F.when(F.col(x).isNotNull(), F.col('wght'))))
).alias(f'wgtd_{x}')
for x in ['var1', 'var2']
]))
轉載請註明出處，本文鏈接：https://www.uj5u.com/yidong/474818.html
標籤：Python 阿帕奇火花 pyspark apache-spark-sql