Search code examples
pythonpysparkisin

pyspark - isin between columns


I am trying to use isin function to check if a value of a pyspark datarame column appears on the same row of another column.

+---+-------------+----+------------+--------+
| ID|         date| loc|   main_list|  GOAL_f|
+---+-------------+----+------------+--------+
|ID1|   2017-07-01|  L1|        [L1]|       1|
|ID1|   2017-07-02|  L1|        [L1]|       1|
|ID1|   2017-07-03|  L2|        [L1]|       0|
|ID1|   2017-07-04|  L2|     [L1,L2]|       1|
|ID1|   2017-07-05|  L1|     [L1,L2]|       1|
|ID1|   2017-07-06|  L3|     [L1,L2]|       0|
|ID1|   2017-07-07|  L3|  [L1,L2,L3]|       1|
+---+-------------+----+------------+--------+

But I am getting errors when trying to collect the main_list for comparison. Here is what I tried unsuccessfully:

df.withColumn('GOAL_f', F.col('loc').isin(F.col('main_list').collect())

Consolidated code:

w = Window.partitionBy('id').orderBy('date').rowsBetween(Window.unboundedPreceeding,-1)
df.withColumn('main_list', F.collect_set('loc').over(w))
  .withColumn('GOAL_f', F.col('loc').isin(F.col('main_list').collect())

Solution

  • You could reverse the query, not asking if value is in something, but if something contains the value.

    Example:

    from pyspark.sql import SparkSession
    import pyspark.sql.functions as F
    
    
    if __name__ == "__main__":
        spark = SparkSession.builder.getOrCreate()
        data = [
            {"loc": "L1", "main_list": ["L1", "L2"]},
            {"loc": "L1", "main_list": ["L2"]},
        ]
        df = spark.createDataFrame(data=data)
        df = df.withColumn(
            "GOAL_f",
            F.when(F.array_contains(F.col("main_list"), F.col("loc")), 1).otherwise(0),
        )
    

    Result:

    +---+---------+------+
    |loc|main_list|GOAL_f|
    +---+---------+------+
    |L1 |[L1, L2] |1     |
    |L1 |[L2]     |0     |
    +---+---------+------+