Conditional logic on groupBy using pyspark-sql-dataframe - group-by

Question

I have a pyspark sql dataframe that looks like this:
id | code
------------------------
| 1 | 02 |
| 1 | 03 |
| 1 | 06 |
| 2 | 02 |
| 2 | 04 |
| 2 | 02 |
| 3 | 06 |
| 3 | 04 |
And I am trying to get an output like this:
id | bin
------------------------
| 1 | 1 |
| 2 | 0 |
| 3 | 1 |
The logic for bin is that if any entry of id contains the code 03 OR 06 then the value of its bin=1 else bin=0. For example, id=1 has bin=1 because some of its code contains 03 and 06; id=2 has bin=0 because none of its code contains either 03 or 06; and id=3 has bin=0 because one its code contains 03.
I have tried using groupBy together with agg but I can only get as far as countDistinct or sum or some flavour of it. Any help will be much appreciated

Answer 1

import pyspark.sql.functions as func
from pyspark.sql import types as T
def bin_fxn(x):
code = set(['03','06'])
return 1 - int(set(x).intersection(code)==set())
def main():
df = sc.parallelize([[1, '02'],
[1, '03'],
[1, '06'],
[2, '02'],
[2, '04'],
[2, '02'],
[3, '06'],
[3, '04']]).toDF(['id', 'code'])
bin_def = func.udf(bin_fxn, T.IntegerType())
df.groupBy('id').agg(bin_def(func.collect_list('code')).alias('bin')).show()
return
The code above achieves the goal. The use of UDF together with pyspark.sql.functions had to be made. For details visit the following website