简体   繁体   中英

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

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

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM