pySpark - строка означает исключение нулей

Question

У меня есть искровой фрейм данных, состоящий из> 1200 столбцов, и мне нужно взять среднее число наборов столбцов по строкам, исключая нулевые значения.

Например, для следующего фрейма данных

id | col1 | col2 | col3
1  |  0   |  2   |  3
2  |  4   |  2   |  3
3  |  1   |  0   |  3
4  |  0   |  0   |  0

Я бы ожидал:

id | mean 
1  |  2.5  
2  |  3  
3  |  2
4  |  0  

В Python я знаю, что можно достичь чего-то подобного, используя стратегию, подобную

data[data == 0] = np.nan
means = np.nanmean(data[:, 1:], axis=1)

Но яя не уверен, как сделать то же самое в pySpark.

Answer 1

Вы можете использовать что-то вроде ниже

>>> import pyspark.sql.functions as F
>>> 
>>> df.show()
+---+----+----+----+
| id|col1|col2|col3|
+---+----+----+----+
|  1|   0|   2|   3|
|  2|   4|   2|   3|
|  3|   1|   0|   3|
|  4|   0|   0|   0|
+---+----+----+----+

>>> cols = [i for i in df.columns if i != 'id']
>>> df = df.withColumn('mean',\
...     sum([df[i] for i in cols])/ \
...     sum([F.when(df[i]>0,1).otherwise(0) for i in cols])). \
...     fillna(0,'mean')
>>> 
>>> df.show()
+---+----+----+----+----+
| id|col1|col2|col3|mean|
+---+----+----+----+----+
|  1|   0|   2|   3| 2.5|
|  2|   4|   2|   3| 3.0|
|  3|   1|   0|   3| 2.0|
|  4|   0|   0|   0| 0.0|
+---+----+----+----+----+