Adding a conditional filter clause

user_2753501 0
Viewed 374

This is my dataframe and filters:

df = spark.range(3)

filter1 = (F.col('id') == 1)
filter2 = (F.col('id') == 2)
flag = False

I want to make filter2 dependent on flag. When flag = True, I want filter2 to take effect and not to take effect when flag = False.

I know I can do

filt = filter1
if flag:
    filt = filt | filter2
df = df.filter(filt)

I wonder if this would be possible in one line, for example utilizing conditional expressions/ ternary operators, etc.

This did not work:

df = df.filter(filter1 | filter2 if flag else False)

Desired result when flag = False:

+---+
| id|
+---+
|  1|
+---+

Desired result when flag = True:

+---+
| id|
+---+
|  1|
|  2|
+---+
pyspark apache-spark-sql apache-spark dataframe conditional-operator
3 Answers

You can use filter2 with a bitwise & condition on the flag . This will only return True when both the filter2 and flag returns True else False. Use this with filter1 with an or condition.

output = df.filter(filter1 |(filter2 & F.lit(flag)))
user_98406370

I tried to nest it inside a when-otherwise block , surprisingly it works in a way -

Data Preparation

sparkDF = sql.range(3)

sparkDF.show()
+---+
| id|
+---+
|  0|
|  1|
|  2|
+---+

True Flag - Filter - When - Otherwise

filter1 = (F.col('id') == 1)
filter2 = (F.col('id') == 2)
flag = True

sparkDF.filter(F.when(F.lit(flag),filter1).otherwise(filter2)).show()
+---+
| id|
+---+
|  1|
+---+

False Flag - Filter - When - Otherwise

filter1 = (F.col('id') == 1)
filter2 = (F.col('id') == 2)
flag = False

sparkDF.filter(F.when(F.lit(flag),filter1).otherwise(filter2)).show()
+---+
| id|
+---+
|  2|
+---+

However I do suspect , if the filter conditions will add multiple columns to filter upon this would break.

Maybe you can have a further nested when-when-*-otherwise to work in that case.

user_91089120

The following worked. But I am still open for suggestions employing python's conditional expressions/ ternary operators. Employing them, conditions would not be repeated - as can be seen, now filter1 is repeated two times.

df = df.filter(F.when(F.lit(flag), filter1 | filter2).otherwise(filter1))
user_27535010
Related