pyspark 计算多列的平均值/总和，忽略空值

Question

我有这样的数据框。

cols=['a','b','c','d','e','f']
find_mean =  F.expr('+'.join(cols))/len(cols)

data=[(
    1,2,3,4,5,None
),(
    1,2,3,4,5,None
),(
    5,4,3,2,1,None
),
(3,4,5,1,2,5)]

schema=T.StructType([
    T.StructField('a',T.IntegerType()),
    T.StructField('b',T.IntegerType()),
    T.StructField('c',T.IntegerType()),
    T.StructField('d',T.IntegerType()),
    T.StructField('e',T.IntegerType()),
    T.StructField('f',T.IntegerType())
])

test=spark.createDataFrame(data,schema)
sum=test.withColumn('sum',find_mean)
sum.display()

创建“平均”列后，我的数据看起来像这样，但有些值为空，因为 f 列中有空值。

a|b|c|d|e|f   |average|
1|2|3|4|5|null|null
1|2|3|4|5|null|null
5|4|3|2|1|null|null
3|4|5|1|2|5   |3.3333

我希望我的数据看起来像这样。

a|b|c|d|e|f   |average|
1|2|3|4|5|null|3
1|2|3|4|5|null|3
5|4|3|2|1|null|3
3|4|5|1|2|5   |3.3333

我想保留“f”列中的空值 - 不想用零填充空值，但在“平均”列中，我希望有有效的数字/小数。

我怎样才能让它发挥作用？

Answer 1

从 Spark 3.1 开始，您可以在计算总和之前过滤数组以删除空值，如下所示：

from pyspark.sql import functions as F

cols = ['a', 'b', 'c', 'd', 'e', 'f']
filtered_array = F.filter(F.array([F.col(i) for i in cols]), lambda c: F.isnotnull(c))
find_mean = F.aggregate(filtered_array, F.lit(0), lambda acc, c: c + acc) / F.size(filtered_array)

这样做，你会得到以下结果：

+---+---+---+---+---+----+------------------+
|a  |b  |c  |d  |e  |f   |sum               |
+---+---+---+---+---+----+------------------+
|1  |2  |3  |4  |5  |NULL|3.0               |
|1  |2  |3  |4  |5  |NULL|3.0               |
|5  |4  |3  |2  |1  |NULL|3.0               |
|3  |4  |5  |1  |2  |5   |3.3333333333333335|
+---+---+---+---+---+----+------------------+