Filtering with conditions in pyspark

Viewed 23

I have this dataframe

df = pandas.DataFrame({'A': [1, 1, 2], 'B': [2, 10, 9], 'C': [3.3, 5.4, 1.5], 'D': [4, 7, 15]},
                      index = ['a1', 'a2', 'a3']) 

I want to do this filtering (if A == 1, then list of columns([B,C]) = 0)

Please anyone know how can I do this in pyspark?

Thank in advance

1 Answers

you can use python's reduce function to create the and chain for the list of columns.

column_list = ['b', 'c']

reduce(lambda x, y: x & y, [func.col(c) == 0 for c in column_list])
# Column<'((b = 0) AND (c = 0))'>

So, you can use the reduce along with condition on column a to filter your data.

(
    (func.col('a') == 1) & 
    reduce(lambda x, y: x & y, [func.col(c) == 0 for c in column_list])
)
# Column<'((a = 1) AND ((b = 0) AND (c = 0)))'>

# use within a `filter()`
data_sdf. \
   filter((func.col('a') == 1) & 
          reduce(lambda x, y: x & y, [func.col(c) == 0 for c in column_list])
          )
Related