PySpark Get Count of True Values in Boolean Columns by Row - pyspark

Question

Please let me know if this is a duplicate. I've been searching all over to try to figure out hwo to do this without writing a user defined function. I have a bunch of boolean columns, each a different quality assurance flag, in a PySpark data frame. All I need to do is create a new column with the number of these columns with a True value, the count of QA checks each row is failing. However, I cannot, for the life of me, figure out an efficient way of doing this. Any ideas, references or links are greatly appreciated!
For instance, for one record with the above columns with the following values...
...I want to create a new column with a count of 2.
Have any good ideas?

Answer 1

Two methods come to mind without using user defined functions .
I'm assuming you have a python list with the boolean column names.
qa_tests = ['qa_flg_xy_equal', 'qa_flg_out_of_bounds_x'] and so forth.
plan a - build in local python a column that is the sum of all boolean columns cast as integers and then put it in spark.
from pyspark.sql.functions import col, lit
from functools import reduce
sum_bools = reduce(lambda acc, v: acc + col(v).cast("integer"), qa_tests, lit(0))
sum_bools is just an automatic way of writing lit(0) + col("qa_flg_xy_equal").cast("integer") + col("qa_flg_out_of_bounds_x").cast("integer") + ...
Here is how sum_bools is defnied:
>> Column<'((0 + CAST(qa_flg_xy_equal AS INT)) + CAST(qa_flg_out_of_bounds_x AS INT))'>
rest of code:
df.withColumn("tests_passed", sum_bools).show(truncate=0)
+---+---------------+----------------------+------------+
|id |qa_flg_xy_equal|qa_flg_out_of_bounds_x|tests_passed|
+---+---------------+----------------------+------------+
|1 |true |false |1 |
|2 |false |false |0 |
|3 |true |true |2 |
+---+---------------+----------------------+------------+
plan b - we can use Array columns to collect all booleans into one array value, filter only the true and check the size of the array after the filter.
No need to keep all the verbose steps below, you can write it in one withColumn of course.
from pyspark.sql.functions import array, filter, size
df \
.withColumn("qa_results", array(*qa_tests)) \
.withColumn("passed_results", filter(
"qa_results",
lambda test: test
)) \
.withColumn("passed_results_count", size("passed_results"))
+---+---------------+----------------------+--------------+--------------+--------------------+
| id|qa_flg_xy_equal|qa_flg_out_of_bounds_x| qa_results|passed_results|passed_results_count|
+---+---------------+----------------------+--------------+--------------+--------------------+
| 1| true| false| [true, false]| [true]| 1|
| 2| false| false|[false, false]| []| 0|
| 3| true| true| [true, true]| [true, true]| 2|
+---+---------------+----------------------+--------------+--------------+--------------------+