Only keep rows with specific condition in PySpark - apache-spark

Question

I'm processing my logs using PySpark. I have two dataframes, one: logs DF is storing search queries and the other one: clicks DF is storing clicked document IDs.
Here is their following structure:
+-------------------+-------+----------------------+
|timestamp |user |query |
+-------------------+-------+----------------------+
|2021-12-01 06:14:38|m96cles|minoration |
|2021-12-01 06:32:54|m96ngro|associés |
|2021-12-01 06:40:40|m96mbeg|cessation |
|2021-12-01 07:02:42|m96ngro|membres de société |
|2021-12-01 07:02:58|m96ngro|cumul |
|2021-12-01 07:07:30|m96rara|cessation |
|2021-12-01 07:09:37|m64nesc|INVF |
|2021-12-01 07:16:14|m83ccat|report didentifiation |
+-------------------+-------+----------------------+
+-------------------+-------+------+
|timestamp |user |doc_id|
+-------------------+-------+------+
|2021-12-01 06:14:42|m96cles|783 |
|2021-12-01 06:33:38|m96ngro|6057 |
|2021-12-01 06:40:52|m96mbeg|1407 |
|2021-12-01 06:49:12|m96mbeg|1414 |
|2021-12-01 06:53:19|m51cstr|15131 |
|2021-12-01 06:53:35|m51cstr|14992 |
|2021-12-01 06:53:55|m51cstr|15093 |
|2021-12-01 06:54:20|m51cstr|15110 |
+-------------------+-------+------+
I merged both dataframes by doing df = logs.unionByName(clicks, allowMissingColumns=True), and sorted it by timestamp.
+-------------------+--------+--------------------+------+
| timestamp| user| query|doc_id|
+-------------------+--------+--------------------+------+
|2022-05-31 20:23:40|ozenfada| null| 7931|
|2022-05-31 21:06:44| m97emou| apnée du sommeil| null|
|2022-05-31 21:28:24| m64lbeh| null| 192|
|2022-05-31 21:29:04| m97emou| null| 3492|
+-------------------+--------+--------------------+------+
The idea is to only keep rows with search queries that lead to clicks. I don't want to keep logs that lead to no clicks on documents. In order to achieve this, I'm trying to look at next rows with the same user and see if they clicked at least on one document within 5 minutes. In the end, I only want to keep rows with search query values.
Here's what I've done so far, I tried to create a boolean column:
df = df.withColumn('valid',
df.user == F.lead('user').over(
W.partitionBy(clicks.user, F.window('timestamp', '5 minutes')).orderBy('timestamp')
)
)
Here's the desired output. Notice how only rows with search queries that lead to clicks (a search query row (query != null) that is followed by "clicks" row(s) (doc_id != null) with the same username) have the true flag. Also, the row with query "rech" lead to a query correction "recherche" and therefore shouldn't be flagged as true.
+-------------------+--------+----------------------------------------+------+-----+
|timestamp |user |query |doc_id|valid|
+-------------------+--------+----------------------------------------+------+-----+
|2022-05-31 18:56:47|m97bcar |exemple |null |false|
|2022-05-31 19:22:40|ozenfada|fort |null |true |
|2022-05-31 19:23:40|ozenfada|null |7931 |false|
|2022-05-31 19:24:09|ozenfada|null |1034 |false|
|2022-05-31 21:06:44|m97emou |apnée du sommeil |null |true |
|2022-05-31 21:07:24|m64lbeh |rech |192 |false|
|2022-05-31 21:07:40|m64lbeh |recherche |null |true |
|2022-05-31 21:08:21|m64lbeh |null |3002 |false|
|2022-05-31 21:11:04|m97emou |null |3492 |false|
+-------------------+--------+----------------------------------------+------+-----+
Any help would be greatly appreciated.

Answer 1

The following will get you all the queries that resulted in a click within 5 minutes. Join on user will work with the extra condition that the timestamp difference between query and click is <=5 minutes. The result printed below is for the sample data provided.
# rename columns to avoid ambiguity
logs = logs.withColumnRenamed('timestamp', 'query_timestamp')
clicks = clicks.withColumnRenamed('timestamp', 'click_timestamp')
clicks = clicks.withColumnRenamed('user', 'click_user')
# join on same username, and if the click is within 5 minutes of the query
time_diff_in_seconds = F.unix_timestamp(clicks['click_timestamp']) - F.unix_timestamp(logs['query_timestamp'])
join_cond = (logs['user']==clicks['click_user']) & \
(time_diff_in_seconds >= 0) & \
(time_diff_in_seconds <= 5*60)
df2 = logs.join(clicks, join_cond, how='left')
# drop all queries that didn't lead to a click
df2 = df2.filter(df2['doc_id'].isNotNull())
# select only the necessary columns
df2 = df2.select('query_timestamp', 'user', 'query').distinct()
df2.show()
+-------------------+-------+----------+
| query_timestamp| user| query|
+-------------------+-------+----------+
|2021-12-01 06:14:38|m96cles|minoration|
|2021-12-01 06:40:40|m96mbeg| cessation|
|2021-12-01 06:32:54|m96ngro| associés|
+-------------------+-------+----------+
Update - To handle misspelled queries
Introduce a column that shows the time difference in seconds between a query and a click. After the join both rows will be retained but the time difference will be larger for the misspelled query. So do an orderBy() on time difference and drop the second row. This can be done with a dropDuplicates('click_timestamp', 'user', 'doc_id').
Let's say there was a search for minor in the 2nd row:
+-------------------+-------+--------------------+
| timestamp| user| query|
+-------------------+-------+--------------------+
|2021-12-01 06:14:38|m96cles| minoration|
|2021-12-01 06:14:39|m96cles| minor|
... and rest of the rows
logs = logs.withColumnRenamed('timestamp', 'query_timestamp')
clicks = clicks.withColumnRenamed('timestamp', 'click_timestamp')
clicks = clicks.withColumnRenamed('user', 'click_user')
time_diff_in_seconds = F.unix_timestamp(clicks['click_timestamp']) - F.unix_timestamp(logs['query_timestamp'])
join_cond = (logs['user']==clicks['click_user']) & \
(time_diff_in_seconds >= 0) & \
(time_diff_in_seconds <= 5*60)
df2 = logs.join(clicks, join_cond, how='left')
df2 = df2.withColumn('time_diff_in_seconds', time_diff_in_seconds)
# ensures if a query leads to multiple clicks then duplicates caused due to left join are dropped
df2 = df2.orderBy('time_diff_in_seconds').dropDuplicates(['query_timestamp', 'user', 'query'])
# keep only the latest query that lead to a click
df2 = df2.orderBy('time_diff_in_seconds').dropDuplicates(['click_timestamp', 'user', 'doc_id'])
df2 = df2.filter(df2['doc_id'].isNotNull())
df2 = df2.select('query_timestamp', 'user', 'query')
df2.show()
+-------------------+-------+---------+
| query_timestamp| user| query|
+-------------------+-------+---------+
|2021-12-01 06:14:39|m96cles| minor|
|2021-12-01 06:32:54|m96ngro| associés|
|2021-12-01 06:40:40|m96mbeg|cessation|
+-------------------+-------+---------+
You may need to test for more complex scenarios, and perhaps modify the code a bit. I think this logic would work, based on the sample data.