Question

假设我有以下由pyspark创建的数据框

id  date         deleted
1   2019-02-07     true
1   2019-02-04     false
2   2019-02-01     true
3   2019-02-08     false
3   2019-02-06     true

我想从最早的日期到现在（例如2019-02-09）每天对该表重新编制索引，并且最早的日期基于每个ID，例如，对于ID 1，最早的日期是2019- 02-04，id 3，最早的日期是2019-02-06。预期结果是：

id  date         deleted
1   2019-02-04     false
1   2019-02-05     null
1   2019-02-06     null
1   2019-02-07     true
1   2019-02-08     null
1   2019-02-09     null

2   2019-02-01     true
2   2019-02-02     null
      ...
2   2019-02-09     null

3   2019-02-06     true
3   2019-02-07     null
3   2019-02-08     false
3   2019-02-09     null

我知道如何基于所有id最早的日期（即2019-02-01），然后构造一个数据框，其中包含每个日期从2019-02-01到2019-02-09的所有日期id（交叉联接），然后左联接原始数据帧。这种方法的问题在于，如果存在一个日期为1980-01-01的数据，则重新索引将为所有id填充1980-01-01到现在的所有数据，这是没有意义的，并且会影响此数据帧上的以下ETL。

最早基于每个分区的日期，都没有找到一种很好的方法。

Answer 1

假设您原来的DataFrame被称为df，并且date列的类型实际上是DateType：

import pyspark.sql.functions as F
from pyspark.sql.types import DateType, ArrayType
import datetime

# create a UDF to create a range of dates from a start
# date until today
def construct_date_range(start_date):
    ndays = (datetime.datetime.today() - start_date).days
    return reversed([base - datetime.timedelta(days=x) for x in range(0, ndays+1)])
date_range_udf = F.udf(construct_date_range, ArrayType(DateType()))

# find the first date for each id, and create a record for
# all dates since the first
id_dates = (
    df
    .groupBy('id')
    .agg(F.min('date').alias('min_date'))
    .withColumn('date_arr', construct_date_range('min_date'))
    .select('id', F.explode('date_arr').alias('date'))
)

result = id_dates.join(df, on=['id','date'], how='left')

Answer 2

基于@abeboparebop的解决方案，我修复了一些格式问题，并使其如下工作：

import pyspark.sql.functions as F
from pyspark.sql.types import DateType, ArrayType
import pandas as pd

from datetime import datetime

import pandas as pd

SYDNEY_TZ = "Australia/Sydney"

def _utc_now():
    return datetime.utcnow()

def _current_datetime_index(timezone=SYDNEY_TZ):
    return pd.DatetimeIndex([_utc_now()]).tz_localize("UTC").tz_convert(timezone).tz_localize(None)


def current_datetime(timezone=SYDNEY_TZ):
    return _current_datetime_index(timezone).to_pydatetime()[0]

def generate_date_list(date_from, date_to=None):
    if date_to is None:
        date_to = current_datetime()
    return pd.date_range(date_from.date(), date_to.date(), freq="D").date.tolist()


def construct_date_range(start_date):
    return generate_date_list(pd.to_datetime(start_date))


date_range_udf = F.udf(construct_date_range, ArrayType(DateType()))


id_dates = (
    given_df
    .groupBy('id')
    .agg(F.min('date').alias('min_date'))
    .withColumn('date_arr', date_range_udf(F.col('min_date')))
    .select('id', F.explode('date_arr').alias('date'))
)

result = id_dates.join(given_df, on=['id', 'date'], how='left')

如何根据每个分区重新索引数据帧

2 个答案: