PySpark и данные временных рядов: как разумно избежать дублирования дат?

Question

У меня есть следующий образец данных Spark

import pandas as pd
import pyspark
import pyspark.sql.functions as fn
from pyspark.sql.window import Window

raw_df = pd.DataFrame([
    (1115, dt.datetime(2019,8,5,18,20), dt.datetime(2019,8,5,18,40)),
    (484, dt.datetime(2019,8,5,18,30), dt.datetime(2019,8,9,18,40)),
    (484, dt.datetime(2019,8,4,18,30), dt.datetime(2019,8,6,18,40)),
    (484, dt.datetime(2019,8,2,18,30), dt.datetime(2019,8,3,18,40)),
    (484, dt.datetime(2019,8,7,18,50), dt.datetime(2019,8,9,18,50)),
    (1115, dt.datetime(2019,8,6,18,20), dt.datetime(2019,8,6,18,40)),
], columns=['server_id', 'start_time', 'end_time'])
df = spark.createDataFrame(raw_df)

что приводит к

+---------+-------------------+-------------------+
|server_id|         start_time|           end_time|
+---------+-------------------+-------------------+
|     1115|2019-08-05 18:20:00|2019-08-05 18:40:00|
|      484|2019-08-05 18:30:00|2019-08-09 18:40:00|
|      484|2019-08-04 18:30:00|2019-08-06 18:40:00|
|      484|2019-08-02 18:30:00|2019-08-03 18:40:00|
|      484|2019-08-07 18:50:00|2019-08-09 18:50:00|
|     1115|2019-08-06 18:20:00|2019-08-06 18:40:00|
+---------+-------------------+-------------------+

Это указывает диапазоны дат использования каждого сервера. Я хочу преобразовать это во временной ряд непересекающихся дат.

Я хотел бы достичь этого без использования UDF.

Это то, что я делаю сейчас, что неправильно

w = Window().orderBy(fn.lit('A'))
# Separate start/end date of usage into rows
df = (df.withColumn('start_end_time', fn.array('start_time', 'end_time'))
    .withColumn('event_dt', fn.explode('start_end_time'))
    .withColumn('row_num', fn.row_number().over(w)))
# Indicate start/end date of the usage (start date will always be on odd rows)
df = (df.withColumn('is_start', fn.when(fn.col('row_num')%2 == 0, 0).otherwise(1))
    .select('server_id', 'event_dt', 'is_start'))

который дает

+---------+-------------------+--------+
|server_id|           event_dt|is_start|
+---------+-------------------+--------+
|     1115|2019-08-05 18:20:00|       1|
|     1115|2019-08-05 18:40:00|       0|
|      484|2019-08-05 18:30:00|       1|
|      484|2019-08-09 18:40:00|       0|
|      484|2019-08-04 18:30:00|       1|
|      484|2019-08-06 18:40:00|       0|
|      484|2019-08-02 18:30:00|       1|
|      484|2019-08-03 18:40:00|       0|
|      484|2019-08-07 18:50:00|       1|
|      484|2019-08-09 18:50:00|       0|
|     1115|2019-08-06 18:20:00|       1|
|     1115|2019-08-06 18:40:00|       0|
+---------+-------------------+--------+

Но конечный результат, которого я хотел бы достичь, заключается в следующем:

+---------+-------------------+--------+
|server_id|           event_dt|is_start|
+---------+-------------------+--------+
|     1115|2019-08-05 18:20:00|       1|
|     1115|2019-08-05 18:40:00|       0|
|     1115|2019-08-06 18:20:00|       1|
|     1115|2019-08-06 18:40:00|       0|
|      484|2019-08-02 18:30:00|       1|
|      484|2019-08-03 18:40:00|       0|
|      484|2019-08-04 18:30:00|       1|
|      484|2019-08-09 18:50:00|       0|
+---------+-------------------+--------+

Таким образом, для server_id 484 У меня есть фактические даты начала и окончания без какого-либо шума между ними.

Есть ли у вас какие-либо предложения о том, как этого добиться без использования UDF?

Спасибо

Answer 1

IIUC, это одна из проблем, которую можно решить с помощью функции Window lag (), sum (), чтобы добавить метку подгруппы для упорядоченных последовательных строк, которые соответствуют некоторым конкретным условиям. Подобно тому, что мы делаем в Пандах, используя shift () + cumsum ().

1. Настройте спецификацию окна w1:

   w1 = Window.partitionBy('server_id').orderBy('start_time')
   

   и рассчитать следующее:

   • max ('end_time'): максимум end_time перед текущей строкой над окном w1
   • лаг ('end_time'): предыдущий end_time
   • sum ('prev_end_time

   Вышеупомянутые три элемента могут соответствовать Pandas cummax (), shift () и cumsum ().

2. Рассчитайте df1, обновив df.end_time с помощью max(end_time).over(w1) и установить метку подгруппы g, а затем сделать groupby(server_id, g) рассчитать min(start_time) а также max(end_time)

   df1 = df.withColumn('end_time', fn.max('end_time').over(w1)) \
           .withColumn('g', fn.sum(fn.when(fn.lag('end_time').over(w1) < fn.col('start_time'),1).otherwise(0)).over(w1)) \
           .groupby('server_id', 'g') \
           .agg(fn.min('start_time').alias('start_time'), fn.max('end_time').alias('end_time'))
   
   df1.show()
   +---------+---+-------------------+-------------------+
   |server_id|  g|         start_time|           end_time|
   +---------+---+-------------------+-------------------+
   |     1115|  0|2019-08-05 18:20:00|2019-08-05 18:40:00|
   |     1115|  1|2019-08-06 18:20:00|2019-08-06 18:40:00|
   |      484|  0|2019-08-02 18:30:00|2019-08-03 18:40:00|
   |      484|  1|2019-08-04 18:30:00|2019-08-09 18:50:00|
   +---------+---+-------------------+-------------------+
   

3. После того, как у нас есть df1, мы можем разделить данные, используя два выбора, а затем объединить набор результатов:

   df_new = df1.selectExpr('server_id', 'start_time as event_dt', '1 as is_start').union(
            df1.selectExpr('server_id', 'end_time as event_dt', '0 as is_start')
   )        
   
   df_new.orderBy('server_id', 'event_dt').show()                                                                            
   +---------+-------------------+--------+
   |server_id|           event_dt|is_start|
   +---------+-------------------+--------+
   |      484|2019-08-02 18:30:00|       1|
   |      484|2019-08-03 18:40:00|       0|
   |      484|2019-08-04 18:30:00|       1|
   |      484|2019-08-09 18:50:00|       0|
   |     1115|2019-08-05 18:20:00|       1|
   |     1115|2019-08-05 18:40:00|       0|
   |     1115|2019-08-06 18:20:00|       1|
   |     1115|2019-08-06 18:40:00|       0|
   +---------+-------------------+--------+