PySpark loop element of column as condtition in where clause used to count row

Agi Nator 0
Viewed 17

The data looks like: ID_device, ID_contract

The scenario, I need a new column filled with 0 or 1 based on the conditon if a ID_contract is assoziate with more than 1 (>) ID_device the column value for this ID_contract is 1 else its 0.

My approach was to iterate over all ID_contract values, take the row value as an condition in a where clause e.g. data.where(data.ID_contract == row value).count() to get the number of ID_devices for each ID_contract. This was done in a udf but it through an error saying its not allowed to perform dataframe calculations within a udf.

Any ideas how to solve this problem?

Great trhanks in advance!

This was my logic approach:

@udf(returnType=IntegerType()) 
def multidevicecontract(contractnr):

   amount_devices = data.where(data.ID_contract == contractnr).count()

   return amount_devices

data = data.withColumn("multidevicecontract", when(multidevicecontract(data.ID_contract) > 1,1).otherwise(0))

The exact error: PicklingError: Could not serialize object: TypeError: cannot pickle '_thread.RLock' object

foreach where-clause count pyspark python
1 Answers

First of all, if you use the UDF in pyspark, the type that you pass into the function is a data type (e.g integer, float etc based on your input column) but not a dataframe and it just passes your column value. Therefore, you can't use any pyspark API inside the UDF.

To achieve your goal, you can just use the window function. Here is the example dataframe:

df = spark.createDataFrame(
    [(idx, randrange(10)) for idx in range(19)],
    schema=['id_device', 'id_contract']
)

df.show(20, False)
+---------+-----------+
|id_device|id_contract|
+---------+-----------+
|0        |7          |
|1        |6          |
|2        |4          |
|3        |2          |
|4        |8          |
|5        |9          |
|6        |7          |
|7        |2          |
|8        |4          |
|9        |8          |
|10       |2          |
|11       |7          |
|12       |1          |
|13       |1          |
|14       |0          |
|15       |5          |
|16       |5          |
|17       |8          |
|18       |0          |
+---------+-----------+

From the above example, only id_contract 6 and 9 associate with 1 id device:

df\
    .withColumn('id_contract_multi', func.when(func.count('id_device').over(Window.partitionBy('id_contract')) > 1, func.lit(1)).otherwise(func.lit(0)))\
    .orderBy('id_device')\
    .show(20, False)
+---------+-----------+-----------------+
|id_device|id_contract|id_contract_multi|
+---------+-----------+-----------------+
|0        |7          |1                |
|1        |6          |0                |
|2        |4          |1                |
|3        |2          |1                |
|4        |8          |1                |
|5        |9          |0                |
|6        |7          |1                |
|7        |2          |1                |
|8        |4          |1                |
|9        |8          |1                |
|10       |2          |1                |
|11       |7          |1                |
|12       |1          |1                |
|13       |1          |1                |
|14       |0          |1                |
|15       |5          |1                |
|16       |5          |1                |
|17       |8          |1                |
|18       |0          |1                |
+---------+-----------+-----------------+
Jonathan Lam0
Related