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