如何用PySpark计算DataFrame各列非零元素的平均值？

现有如下PySpark代码：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame([(0.0, 1.2, -1.3), (0.0, 0.0, 0.0),
                            (-17.2, 20.3, 15.2), (23.4, 1.4, 0.0),], 
                            ['col1', 'col2', 'col3'])

df1  = df.agg(F.avg('col1'))
df2  = df.agg(F.avg('col2'))
df3  = df.agg(F.avg('col3'))

对应的DataFrame数据：

ID COL1 COL2 COL3  
1  0.0     1.2    -1.3  
2  0.0     0.0     0.0  
3 -17.2   20.3    15.2
4  23.4   1.4     0.0

当前代码会计算包含零值的各列平均值，结果为1.54、5.725、3.47，但需要计算仅非零元素的平均值，期望结果如下：

avg1 avg2 avg3
1   3.1  7.6  6.9

请问该如何实现？

可以通过F.when()函数过滤掉列中的0值，再对过滤后的非零值计算平均值。PySpark的avg()函数会自动忽略null值，因此把0值替换为null后，就能得到仅非零元素的平均值。

实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame([(0.0, 1.2, -1.3), (0.0, 0.0, 0.0),
                            (-17.2, 20.3, 15.2), (23.4, 1.4, 0.0),], 
                            ['col1', 'col2', 'col3'])

# 一次性计算所有列的非零平均值，并重命名结果列
result_df = df.agg(
    F.avg(F.when(F.col('col1') != 0, F.col('col1'))).alias('avg1'),
    F.avg(F.when(F.col('col2') != 0, F.col('col2'))).alias('avg2'),
    F.avg(F.when(F.col('col3') != 0, F.col('col3'))).alias('avg3')
)

result_df.show()

代码说明 #

• F.when(F.col('col1') != 0, F.col('col1'))：当列值不等于0时保留原值，否则返回null
• F.avg(...)：自动忽略null值（即原0值），计算非零元素的平均值
• alias('avg1')：将结果列重命名为期望的名称，方便识别

执行后得到的结果：

+----+----+----+
|avg1|avg2|avg3|
+----+----+----+
| 3.1| 7.6| 6.9|
+----+----+----+

另一种实现方式 #

也可以通过非零元素总和除以非零元素个数的方式计算，结果一致：

result_df = df.agg(
    (F.sum(F.when(F.col('col1') != 0, F.col('col1'))) / F.count(F.when(F.col('col1') != 0, True))).alias('avg1'),
    (F.sum(F.when(F.col('col2') != 0, F.col('col2'))) / F.count(F.when(F.col('col2') != 0, True))).alias('avg2'),
    (F.sum(F.when(F.col('col3') != 0, F.col('col3'))) / F.count(F.when(F.col('col3') != 0, True))).alias('avg3')
)

内容的提问来源于stack exchange，提问作者alryosha