Create flag if ID maps to multiple other IDs within a window

user_6029749 0
Viewed 116

I am trying to create a flag that shows which id1 values have a duplicate mapping to id2 over an arbitrary window n (n = 2 in the example below). For example, given this DataFrame:

df = spark.createDataFrame(
[("2010-03-10", "A", "X"),
 ("2010-03-10", "A", "Y"),
 ("2010-03-10", "B", "Z"),
 ("2010-04-10", "A", "X"),
 ("2010-04-10", "A", "Y"),
 ("2010-04-10", "B", "Z"),
 ("2010-05-10", "A", "X"),
 ("2010-05-10", "A", "Y"),
 ("2010-05-10", "B", "Z"),
 ("2010-06-10", "A", "X"),
 ("2010-06-10", "B", "Z"),
 ("2010-07-10", "A", "X"),
 ("2010-07-10", "B", "Z")],
 ("date", "id1", "id2")
)
df.show()

+----------+---+---+
|      date|id1|id2|
+----------+---+---+
|2010-03-10|  A|  X|
|2010-03-10|  A|  Y|
|2010-03-10|  B|  Z|
|2010-04-10|  A|  X|
|2010-04-10|  A|  Y|
|2010-04-10|  B|  Z|
|2010-05-10|  A|  X|
|2010-05-10|  A|  Y|
|2010-05-10|  B|  Z|
|2010-06-10|  A|  X|
|2010-06-10|  B|  Z|
|2010-07-10|  A|  X|
|2010-07-10|  B|  Z|
+----------+---+---+

With expected output:

+----------+---+---+--------+
|      date|id1|id2|dup_flag|
+----------+---+---+--------+
|2010-03-10|  A|  X|       1|
|2010-03-10|  A|  Y|       1|
|2010-03-10|  B|  Z|       0|
|2010-04-10|  A|  X|       1|
|2010-04-10|  A|  Y|       1|
|2010-04-10|  B|  Z|       0|
|2010-05-10|  A|  X|       1|
|2010-05-10|  A|  Y|       1|
|2010-05-10|  B|  Z|       0|
|2010-06-10|  A|  X|       1|
|2010-06-10|  B|  Z|       0|
|2010-07-10|  A|  X|       0|
|2010-07-10|  B|  Z|       0|
+----------+---+---+--------+

Where dup_flag is a flag that indicates whether or not id1 contains a duplicate mapping to id2 somewhere in the past 2 months.

That is, for id1 B, it never maps to any other id2 than Z, hence it is never flagged as a duplicate. For id1 A we do have duplicate mappings. In this example I put the window size to 2, which means that if A maps to multiple id2 within any date of a window that goes at most one month back (i.e. the current month as well as the past month should be included in the window), it should receive a flag. Hence, A receives this flag up until 2010-07-10, as in this period, we have a window that consists of both 2010-06-10 and 2010-07-10, in which A does not map to multiple id2 (it only maps to id2 X in this window).

pyspark apache-spark
1 Answers

A somewhat hacky way to solve this. Essentially you create a window which is range-based instead of row-based, and do a count distinct over that window. Somehow Spark doesn't support the latter, so I needed to hack around it by using size(collect_set) which is essentially the same as count distinct. Also Spark doesn't support range-based windows ordered by dates, so I needed to cast the column into integer type and use a magic number 2678400 which is equal to the number of seconds in 31 days.

import pyspark.sql.functions as F
from pyspark.sql.window import Window

### initialize df as in the question
### df = ...

df = df.withColumn('date', F.col('date').cast('timestamp').cast('bigint'))
flag = (F.size(F.collect_set(F.col('id2')).over(Window.partitionBy('id1').orderBy('date').rangeBetween(-2678400, 0))) > 1).cast('int')
df = df.withColumn('flag', flag)
df = df.withColumn('date', F.col('date').cast('timestamp').cast('date'))

The idea being something like

COUNT DISTINCT id2 OVER (PARTITION BY id1 ORDER BY date RANGE BETWEEN INTERVAL '1' MONTH PRECEDING AND CURRENT ROW)
user_141657300
Related