Check multiple columns for any column greater than zero using a regex

user_11651146 0
Viewed 1104

I need to apply a when function on multiple columns. I want to check if at least one of the columns has a value greater than 0.

This is my solution:

df.withColumn("any value", F.when(
   (col("col1") > 0) |
   (col("col2") > 0) |
   (col("col3") > 0) |
   ...
   (col("colX") > 0)
   , "any greater than 0").otherwise(None))

Is it possible to do the same task with a regex, so I don't have to write all the column names?

pyspark apache-spark-sql apache-spark
2 Answers

So let's create sample data:

 df = spark.createDataFrame(
    [(0, 0, 0, 0), (0, 0, 2, 0), (0, 0, 0, 0), (1, 0, 0, 0)],
    ['a', 'b', 'c', 'd']
)

Then, you can build your condition from a list of columns (say all the columns of the dataframe) using map and reduce like this:

cols = df.columns
from pyspark.sql import functions as F
condition = reduce(lambda a, b: a | b, map(lambda c: F.col(c) > 0, cols))
df.withColumn("any value", F.when(condition, "any greater than 0")).show()

which yields:

+---+---+---+---+------------------+
|  a|  b|  c|  d|         any value|
+---+---+---+---+------------------+
|  0|  0|  0|  0|              null|
|  0|  0|  2|  0|any greater than 0|
|  0|  0|  0|  0|              null|
|  1|  0|  0|  0|any greater than 0|
+---+---+---+---+------------------+
user_88936860

Another way you could have this done is create an array, use forall to check and conditionally assign values. Code below

df = df.withColumn('any value', array(df.columns)).withColumn('any value',when(forall('any value',lambda x: x==0),None).otherwise("any greater than 0"))

df.show()


+---+---+---+---+------------------+
|  a|  b|  c|  d|         any value|
+---+---+---+---+------------------+
|  0|  0|  0|  0|              null|
|  0|  0|  2|  0|any greater than 0|
|  0|  0|  0|  0|              null|
|  1|  0|  0|  0|any greater than 0|
+---+---+---+---+------------------+
user_89869750
Related