How to add a new column to pySpark dataframe which contains count its column values which are greater to 0?

Question

I want to add a new column to pyspark dataframe which contains count of all columns values which are greater to 0 in a particular row.

Here is my demo dataframe.

+-----------+----+----+----+----+----+----+
|customer_id|2010|2011|2012|2013|2014|2015|
+-----------+----+----+----+----+----+----+
|     1     |  0 |  4 |  0 | 32 |  0 | 87 |
|     2     |  5 |  5 | 56 | 23 |  0 | 09 |
|     3     |  6 |  6 | 87 |  0 | 45 | 23 |
|     4     |  7 |  0 | 12 | 89 | 78 | 0  |
|     6     |  0 |  0 |  0 | 23 | 45 | 64 |
+-----------+----+----+----+----+----+----+

Above data frame have visit by a customer in a year. I want to count how many years a customer visited. So i need a column visit_count which is having count of visits in year (2010,2011,2012,2013,2014,2015) having value greater to 0.

+-----------+----+----+----+----+----+----+-----------+
|customer_id|2010|2011|2012|2013|2014|2015|visit_count|
+-----------+----+----+----+----+----+----+-----------+
|     1     |  0 |  4 |  0 | 32 |  0 | 87 |    3      |
|     2     |  5 |  5 | 56 | 23 |  0 | 09 |    5      |
|     3     |  6 |  6 | 87 |  0 | 45 | 23 |    5      |
|     4     |  7 |  0 | 12 | 89 | 78 | 0  |    4      |
|     6     |  0 |  0 |  0 | 23 | 45 | 64 |    3      |
+-----------+----+----+----+----+----+----+-----------+

How to Achieve this?

Answer 1

Try this:

df.withColumn('visit_count', sum((df[col] > 0).cast('integer') for col in df.columns))