Функция окна Spark SQL со сложным условием

Это, вероятно, проще всего объяснить на примере. Предположим, у меня есть DataFrame для входа пользователя на сайт, например:

scala> df.show(5)
+----------------+----------+
|       user_name|login_date|
+----------------+----------+
|SirChillingtonIV|2012-01-04|
|Booooooo99900098|2012-01-04|
|Booooooo99900098|2012-01-06|
|  OprahWinfreyJr|2012-01-10|
|SirChillingtonIV|2012-01-11|
+----------------+----------+
only showing top 5 rows

Я хотел бы добавить к этому столбец, указывающий, когда они стали активным пользователем на сайте. Но есть одно предостережение: есть период времени, в течение которого пользователь считается активным, и по истечении этого периода, если они снова became_active в систему, они будут became_active. Предположим, что этот период составляет 5 дней. Тогда желаемая таблица, полученная из приведенной выше таблицы, будет примерно такой:

+----------------+----------+-------------+
|       user_name|login_date|became_active|
+----------------+----------+-------------+
|SirChillingtonIV|2012-01-04|   2012-01-04|
|Booooooo99900098|2012-01-04|   2012-01-04|
|Booooooo99900098|2012-01-06|   2012-01-04|
|  OprahWinfreyJr|2012-01-10|   2012-01-10|
|SirChillingtonIV|2012-01-11|   2012-01-11|
+----------------+----------+-------------+

Так, в частности, SirChillingtonIV стала became_active дата была сброшена, потому что их второй логин пришел после истечения активного периода, но Booooooo99900098 became_active date не был сброшен во второй раз, когда он/она вошел в систему, потому что он попал в активный период.

Моя первоначальная мысль заключалась в том, чтобы использовать функции окна с lag, а затем использовать значения lag для заполнения столбца became_active; например, что-то началось примерно так:

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions._

val window = Window.partitionBy("user_name").orderBy("login_date")
val df2 = df.withColumn("tmp", lag("login_date", 1).over(window))

Тогда правилом, чтобы заполнить дату became_active было бы, если tmp равно null (т. login_date - tmp >= 5 Если это первый когда-либо логин), или если login_date - tmp >= 5 то became_active = login_date; в противном случае перейдите к следующему последнему значению в tmp и примените одно и то же правило. Это говорит о рекурсивном подходе, в котором я испытываю трудности с созданием способа реализации.

Мои вопросы: Является ли это жизнеспособным подходом, и если да, то как я могу "вернуться" и посмотреть на более ранние значения tmp пока не найду ту, где останавливаюсь? Я не могу, насколько мне известно, перебирать значения Column Spark SQL. Есть ли другой способ достичь этого результата?

Ответ 1

Вот трюк. Импортируйте кучу функций:

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{coalesce, datediff, lag, lit, min, sum}

Определить окна:

val userWindow = Window.partitionBy("user_name").orderBy("login_date")
val userSessionWindow = Window.partitionBy("user_name", "session")

Найдите точки, где начинаются новые сеансы:

val newSession =  (coalesce(
  datediff($"login_date", lag($"login_date", 1).over(userWindow)),
  lit(0)
) > 5).cast("bigint")

val sessionized = df.withColumn("session", sum(newSession).over(userWindow))

Найдите самую раннюю дату в сеансе:

val result = sessionized
  .withColumn("became_active", min($"login_date").over(userSessionWindow))
  .drop("session")

С помощью набора данных, определенного как:

val df = Seq(
  ("SirChillingtonIV", "2012-01-04"), ("Booooooo99900098", "2012-01-04"),
  ("Booooooo99900098", "2012-01-06"), ("OprahWinfreyJr", "2012-01-10"), 
  ("SirChillingtonIV", "2012-01-11"), ("SirChillingtonIV", "2012-01-14"),
  ("SirChillingtonIV", "2012-08-11")
).toDF("user_name", "login_date")

Результат:

+----------------+----------+-------------+
|       user_name|login_date|became_active|
+----------------+----------+-------------+
|  OprahWinfreyJr|2012-01-10|   2012-01-10|
|SirChillingtonIV|2012-01-04|   2012-01-04| <- The first session for user
|SirChillingtonIV|2012-01-11|   2012-01-11| <- The second session for user
|SirChillingtonIV|2012-01-14|   2012-01-11| 
|SirChillingtonIV|2012-08-11|   2012-08-11| <- The third session for user
|Booooooo99900098|2012-01-04|   2012-01-04|
|Booooooo99900098|2012-01-06|   2012-01-04|
+----------------+----------+-------------+

Ответ 2

Рефакторинг приведенного выше ответа для работы с Pyspark

В Pyspark вы можете сделать, как Pyspark ниже.

create data frame

df = sqlContext.createDataFrame(
[
("SirChillingtonIV", "2012-01-04"), 
("Booooooo99900098", "2012-01-04"), 
("Booooooo99900098", "2012-01-06"), 
("OprahWinfreyJr", "2012-01-10"), 
("SirChillingtonIV", "2012-01-11"), 
("SirChillingtonIV", "2012-01-14"), 
("SirChillingtonIV", "2012-08-11")
], 
("user_name", "login_date"))

Приведенный выше код создает фрейм данных, как показано ниже

+----------------+----------+
|       user_name|login_date|
+----------------+----------+
|SirChillingtonIV|2012-01-04|
|Booooooo99900098|2012-01-04|
|Booooooo99900098|2012-01-06|
|  OprahWinfreyJr|2012-01-10|
|SirChillingtonIV|2012-01-11|
|SirChillingtonIV|2012-01-14|
|SirChillingtonIV|2012-08-11|
+----------------+----------+

Теперь мы хотим сначала выяснить разницу между login_date более 5 дней.

Для этого сделайте как ниже.

Необходимый импорт

from pyspark.sql import functions as f
from pyspark.sql import Window


# defining window partitions  
login_window = Window.partitionBy("user_name").orderBy("login_date")
session_window = Window.partitionBy("user_name", "session")

session_df = df.withColumn("session", f.sum((f.coalesce(f.datediff("login_date", f.lag("login_date", 1).over(login_window)), f.lit(0)) > 5).cast("int")).over(login_window))

Когда мы запустим приведенную выше строку кода, если date_diff равен NULL тогда функция coalesce заменит NULL на 0.

+----------------+----------+-------+
|       user_name|login_date|session|
+----------------+----------+-------+
|  OprahWinfreyJr|2012-01-10|      0|
|SirChillingtonIV|2012-01-04|      0|
|SirChillingtonIV|2012-01-11|      1|
|SirChillingtonIV|2012-01-14|      1|
|SirChillingtonIV|2012-08-11|      2|
|Booooooo99900098|2012-01-04|      0|
|Booooooo99900098|2012-01-06|      0|
+----------------+----------+-------+


# add became_active column by finding the 'min login_date' for each window partitionBy 'user_name' and 'session' created in above step
final_df = session_df.withColumn("became_active", f.min("login_date").over(session_window)).drop("session")

+----------------+----------+-------------+
|       user_name|login_date|became_active|
+----------------+----------+-------------+
|  OprahWinfreyJr|2012-01-10|   2012-01-10|
|SirChillingtonIV|2012-01-04|   2012-01-04|
|SirChillingtonIV|2012-01-11|   2012-01-11|
|SirChillingtonIV|2012-01-14|   2012-01-11|
|SirChillingtonIV|2012-08-11|   2012-08-11|
|Booooooo99900098|2012-01-04|   2012-01-04|
|Booooooo99900098|2012-01-06|   2012-01-04|
+----------------+----------+-------------+