pyspark - 在列之间

Question

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())

Answer 1

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     |
+---+---------+------+