Pyspark - подсчет ненулевых столбцов в кадре данных искры для каждой строки.

Question

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

ID COL1 COL2 COL3
1  0    1    -1 
2  0    0     0 
3 -17   20    15
4  23   1     0

Ожидаемый результат:

ID COL1 COL2 COL3 Count
1   0    1    -1    2
2   0    0     0    0
3  -17   20    15   3
4   23   1     0    1

Answer 1

Существуют различные подходы для достижения этой цели, ниже я перечисляю один из простых подходов -

df = sqlContext.createDataFrame([
    [1,  0,    1,    -1], 
    [2,  0,    0,     0],
    [3, -17,   20,    15],
    [4,  23,   1,     0]], 
    ["ID", "COL1", "COL2", "COL3"]
)

#Check columns list removing ID columns
df.columns[1:]
['COL1', 'COL2', 'COL3']

#import functions
from pyspark.sql import functions as F

#Adding new column count having sum/addition(if column !=0 then 1 else 0)
df.withColumn(
    "count",
    sum([
            F.when(F.col(cl) != 0, 1).otherwise(0) for cl in df.columns[1:]
    ])
).show()


+---+----+----+----+-----+
| ID|COL1|COL2|COL3|count|
+---+----+----+----+-----+
|  1|   0|   1|  -1|    2|
|  2|   0|   0|   0|    0|
|  3| -17|  20|  15|    3|
|  4|  23|   1|   0|    2|
+---+----+----+----+-----+