从两个日期列填充按月数据框

时间:2019-08-28 19:13:05

标签: python pyspark

我有一个这样的PySpark数据框,

+----------+--------+----------+----------+
|id_       | p      |d1        |  d2      |
+----------+--------+----------+----------+
|  1       | A      |2018-09-26|2018-10-26|
|  2       | B      |2018-06-21|2018-07-19|
|  2       | B      |2018-08-13|2018-10-07|
|  2       | B      |2018-12-31|2019-02-27|
|  2       | B      |2019-05-28|2019-06-25|
|  3       |C       |2018-06-15|2018-07-13|
|  3       |C       |2018-08-15|2018-10-09|
|  3       |C       |2018-12-03|2019-03-12|
|  3       |C       |2019-05-10|2019-06-07|
| 4        | A      |2019-01-30|2019-03-01|
| 4        | B      |2019-05-30|2019-07-25|
| 5        |C       |2018-09-19|2018-10-17|
-------------------------------------------

我必须从此数据帧派生具有n列的另一个数据帧。其中每列是从month(min(d1))month(max(d2))的一个月。

我想要在派生数据框中找到一个实际数据框中的行,并且列值必须是该月中的天数。

例如,

对于第一行,其中id_1,而pA,我想在派生数据框中获得一行,其中201809的列值5和列201810的值26

对于id_2pB的第二行,我想在派生数据帧中获得一行,其中201806的列应是9,而201807应该是19

对于倒数第二行,我希望用值201905填充列1,用值201906填充列30,用{{1}填充201907 }}。

因此,基本上,我希望以如下方式填充数据框:对于原始数据框中的每一行,我在派生数据框中都有一行,其中应填充表中与月份对应的列,范围min(d1)到max(d2),其值表示该特定月份的天数。

我目前正在艰难地这样做。我正在制作n列,其中列的范围是从25开始的日期。我用min(d1) to max(d2)填充这些列,然后融合数据并基于1进行过滤。最后汇总此数据框以获得我想要的结果,然后选择最大值value

在代码中,

p

此代码需要大量时间才能在适当的配置下运行。我该如何改善??

如何以更优化的方式完成此任务?使范围内的每个日期似乎都不是最好的解决方案。

所需输出的一小部分显示如下,

d = df.select(F.min('d1').alias('d1'), F.max('d2').alias('d2')).first()

cols = [ c.strftime('%Y-%m-%d') for c in pd.period_range(d.d1, d.d2, freq='D') ]

result = df.select('id_', 'p', *[ F.when((df.d1 <= c)&(df.d2 >= c), 1).otherwise(0).alias(c) for c in cols ])

melted_data = melt(result, id_vars=['id_','p'], value_vars=cols)

melted_data = melted_data.withColumn('Month', F.substring(F.regexp_replace('variable', '-', ''), 1, 6))

melted_data = melted_data.groupBy('id_', 'Month', 'p').agg(F.sum('value').alias('days'))

melted_data = melted_data.orderBy('id_', 'Month', 'days', ascending=[False, False, False])

final_data = melted_data.groupBy('id_', 'Month').agg(F.first('p').alias('p'))

3 个答案:

答案 0 :(得分:3)

我认为它由于freq='D'和数据集上的多个转换而变慢了。

请尝试以下操作:

编辑1:该季度的更新

编辑2:根据评论,开始日期应包含在最终结果中

编辑3:按评论,每日更新

  1. 准备的数据
#Imports
import pyspark.sql.functions as f
from pyspark.sql.functions import when
import pandas as pd

 df.show()
+---+---+----------+----------+
| id|  p|        d1|        d2|
+---+---+----------+----------+
|  1|  A|2018-09-26|2018-10-26|
|  2|  B|2018-06-21|2018-07-19|
|  2|  B|2018-08-13|2018-10-07|
|  2|  B|2018-12-31|2019-02-27|
|  2|  B|2019-05-28|2019-06-25|
|  3|  C|2018-06-15|2018-07-13|
|  3|  C|2018-08-15|2018-10-09|
|  3|  C|2018-12-03|2019-03-12|
|  3|  C|2019-05-10|2019-06-07|
|  4|  A|2019-01-30|2019-03-01|
|  4|  B|2019-05-30|2019-07-25|
|  5|  C|2018-09-19|2018-10-17|
|  5|  C|2019-05-16|2019-05-29| # --> Same month case
+---+---+----------+----------+
  1. 从频率为freq='M'的数据集中获取最小和最大日期
d = df.select(f.min('d1').alias('min'), f.max('d2').alias('max')).first()
dates = pd.period_range(d.min, d.max, freq='M').strftime("%Y%m").tolist()
dates
['201806', '201807', '201808', '201809', '201810', '201811', '201812', '201901', '201902', '201903', '201904', '201905', '201906', '201907']

  1. 现在,使用火花日期运算符和函数的最终商务逻辑
df1 = df.select('id', 
    'p', 
    'd1',
    'd2', *[ (when( (f.trunc(df.d1, "month") == f.trunc(df.d2, "month")) & (f.to_date(f.lit(c),'yyyyMM') == f.trunc(df.d1, "month"))
                        , f.datediff(df.d2 , df.d1) +1 ) # Same month ((Last day - First dat) + 1
        .when(f.to_date(f.lit(c),'yyyyMM') == f.trunc(df.d1, "month") , 
                        f.datediff(f.last_day(f.to_date(f.lit(c),'yyyyMM')) , df.d1) +1 ) # d1 date (Last day - current day)
        .when(f.to_date(f.lit(c),'yyyyMM') == f.trunc(df.d2, "month") , 
                    f.datediff(df.d2, f.to_date(f.lit(c),'yyyyMM')) +1 ) # d2 date (Currentday - Firstday) 
        .when(f.to_date(f.lit(c),'yyyyMM').between(f.trunc(df.d1, "month"), df.d2), 
                    f.dayofmonth(f.last_day(f.to_date(f.lit(c),'yyyyMM')))) # Between date (Total days in month)
        ).otherwise(0) # Rest of the months (0)
    .alias(c) for c in dates ])

df1.show()
+---+---+----------+----------+------+------+------+------+------+------+------+------+------+------+------+------+------+------+
| id|  p|        d1|        d2|201806|201807|201808|201809|201810|201811|201812|201901|201902|201903|201904|201905|201906|201907|
+---+---+----------+----------+------+------+------+------+------+------+------+------+------+------+------+------+------+------+
|  1|  A|2018-09-26|2018-10-26|     0|     0|     0|     5|    26|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  2|  B|2018-06-21|2018-07-19|    10|    19|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  2|  B|2018-08-13|2018-10-07|     0|     0|    19|    30|     7|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  2|  B|2018-12-31|2019-02-27|     0|     0|     0|     0|     0|     0|     1|    31|    27|     0|     0|     0|     0|     0|
|  2|  B|2019-05-28|2019-06-25|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     4|    25|     0|
|  3|  C|2018-06-15|2018-07-13|    16|    13|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  3|  C|2018-08-15|2018-10-09|     0|     0|    17|    30|     9|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  3|  C|2018-12-03|2019-03-12|     0|     0|     0|     0|     0|     0|    29|    31|    28|    12|     0|     0|     0|     0|
|  3|  C|2019-05-10|2019-06-07|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|    22|     7|     0|
|  4|  A|2019-01-30|2019-03-01|     0|     0|     0|     0|     0|     0|     0|     2|    28|     1|     0|     0|     0|     0|
|  4|  B|2019-05-30|2019-07-25|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     2|    30|    25|
|  5|  C|2018-09-19|2018-10-17|     0|     0|     0|    12|    17|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  5|  C|2019-05-16|2019-05-29|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|    14|     0|     0|
+---+---+----------+----------+------+------+------+------+------+------+------+------+------+------+------+------+------+------+

编辑2:更新季度范围:

注意: :从 @jxc的答案中获取季度日期范围字典。我们对这里的最佳解决方案更感兴趣。 @jxc做得非常好,除非有性能问题,否则不要重新设计轮子。

创建日期范围字典:

q_dates = dict([
    (str(c), [ c.to_timestamp().strftime("%Y-%m-%d") ,(c.to_timestamp() + pd.tseries.offsets.QuarterEnd()).strftime("%Y-%m-%d")
     ]) for c in pd.period_range(d.min, d.max, freq='Q')
])


# {'2018Q2': ['2018-04-01', '2018-06-30'], 
# '2018Q3': ['2018-07-01', '2018-09-30'], 
# '2018Q4': ['2018-10-01', '2018-12-31'], 
# '2019Q1': ['2019-01-01', '2019-03-31'], 
# '2019Q2': ['2019-04-01', '2019-06-30'], 
# '2019Q3': ['2019-07-01', '2019-09-30']}

现在在每个季度应用业务逻辑。

df1 = df.select('id', 
    'p', 
    'd1',
    'd2', 
    *[(when( (df.d1.between(q_dates[c][0], q_dates[c][1])) & (f.trunc(df.d1, "month") == f.trunc(df.d2, "month")), 
        f.datediff(df.d2 , df.d1) +1 ) # Same month ((Last day - start day) +1 )
    .when(df.d1.between(q_dates[c][0], q_dates[c][1]), 
        f.datediff(f.to_date(f.lit(q_dates[c][1])), df.d1) +1) # Min date , remaining days (Last day of quarter - Min day)
    .when(df.d2.between(q_dates[c][0], q_dates[c][1]), 
        f.datediff(df.d2, f.to_date(f.lit(q_dates[c][0]))) +1 ) # Max date , remaining days (Max day - Start day of quarter )
    .when(f.to_date(f.lit(q_dates[c][0])).between(df.d1, df.d2), 
        f.datediff(f.to_date(f.lit(q_dates[c][1])), f.to_date(f.lit(q_dates[c][0]))) +1) # All remaining days
    ).otherwise(0)
    .alias(c) for c in q_dates ])

df1.show()
+---+---+----------+----------+------+------+------+------+------+------+
| id|  p|        d1|        d2|2018Q2|2018Q3|2018Q4|2019Q1|2019Q2|2019Q3|
+---+---+----------+----------+------+------+------+------+------+------+
|  1|  A|2018-09-26|2018-10-26|     0|     5|    26|     0|     0|     0|
|  2|  B|2018-06-21|2018-07-19|    10|    19|     0|     0|     0|     0|
|  2|  B|2018-08-13|2018-10-07|     0|    49|     7|     0|     0|     0|
|  2|  B|2018-12-31|2019-02-27|     0|     0|     1|    58|     0|     0|
|  2|  B|2019-05-28|2019-06-25|     0|     0|     0|     0|    34|     0|
|  3|  C|2018-06-15|2018-07-13|    16|    13|     0|     0|     0|     0|
|  3|  C|2018-08-15|2018-10-09|     0|    47|     9|     0|     0|     0|
|  3|  C|2018-12-03|2019-03-12|     0|     0|    29|    71|     0|     0|
|  3|  C|2019-05-10|2019-06-07|     0|     0|     0|     0|    52|     0|
|  4|  A|2019-01-30|2019-03-01|     0|     0|     0|    61|     0|     0|
|  4|  B|2019-05-30|2019-07-25|     0|     0|     0|     0|    32|    25|
|  5|  C|2018-09-19|2018-10-17|     0|    12|    17|     0|     0|     0|
|  5|  C|2019-05-16|2019-05-29|     0|     0|     0|     0|    14|     0|
+---+---+----------+----------+------+------+------+------+------+------+

编辑3:按评论,每日更新

由于此处的评估更多,因此在性能方面需要谨慎。

方法1:数据框/数据集

yyyy-MM-dd格式但以字符串形式获取日期列表

df_dates = pd.period_range(d.min, d.max, freq='D').strftime("%Y-%m-%d").tolist() 

现在,业务逻辑非常简单。它是1或0

df1 = df.select('id'
    , 'p'
    , 'd1'
    ,'d2'
    , *[ (when(f.lit(c).between (df.d1, df.d2),1)) # For date rabge 1
        .otherwise(0) # For rest of days
        .alias(c) for c in df_dates ])

df1.show()

+---+---+----------+----------+----------+----------+----------+
| id|  p|        d1|        d2|2018-06-15|2018-06-16|2018-06-17| # and so on....
+---+---+----------+----------+----------+----------+----------+
|  1|  A|2018-09-26|2018-10-26|         0|         0|         0|
|  2|  B|2018-06-21|2018-07-19|         0|         0|         0|
|  2|  B|2018-08-13|2018-10-07|         0|         0|         0|
|  2|  B|2018-12-31|2019-02-27|         0|         0|         0|
|  2|  B|2019-05-28|2019-06-25|         0|         0|         0|
|  3|  C|2018-06-15|2018-07-13|         1|         1|         1|
|  3|  C|2018-08-15|2018-10-09|         0|         0|         0|
|  3|  C|2018-12-03|2019-03-12|         0|         0|         0|
|  3|  C|2019-05-10|2019-06-07|         0|         0|         0|
|  4|  A|2019-01-30|2019-03-01|         0|         0|         0|
|  4|  B|2019-05-30|2019-07-25|         0|         0|         0|
|  5|  C|2018-09-19|2018-10-17|         0|         0|         0|
|  5|  C|2019-05-16|2019-05-29|         0|         0|         0|
+---+---+----------+----------+----------+----------+----------+
# Due to answer character limit unable to give the result.

方法2:RDD评估

将日期列表作为date object

rdd_dates = [ c.to_timestamp().date() for c in pd.period_range(d.min, d.max, freq='D') ]

使用map中的rdd


df1 = df \
.rdd \
.map(lambda x : tuple([x.id, x.p, x.d1, x.d2 , *[ 1 if ( x.d1 <= c <=x.d2) else  0 for c in rdd_dates]])) \
.toDF(df.columns + [ c.strftime("%Y-%m-%d") for c in rdd_dates])

df1.show()

+---+---+----------+----------+----------+----------+----------+
| id|  p|        d1|        d2|2018-06-15|2018-06-16|2018-06-17| # and so on....
+---+---+----------+----------+----------+----------+----------+
|  1|  A|2018-09-26|2018-10-26|         0|         0|         0|
|  2|  B|2018-06-21|2018-07-19|         0|         0|         0|
|  2|  B|2018-08-13|2018-10-07|         0|         0|         0|
|  2|  B|2018-12-31|2019-02-27|         0|         0|         0|
|  2|  B|2019-05-28|2019-06-25|         0|         0|         0|
|  3|  C|2018-06-15|2018-07-13|         1|         1|         1|
|  3|  C|2018-08-15|2018-10-09|         0|         0|         0|
|  3|  C|2018-12-03|2019-03-12|         0|         0|         0|
|  3|  C|2019-05-10|2019-06-07|         0|         0|         0|
|  4|  A|2019-01-30|2019-03-01|         0|         0|         0|
|  4|  B|2019-05-30|2019-07-25|         0|         0|         0|
|  5|  C|2018-09-19|2018-10-17|         0|         0|         0|
|  5|  C|2019-05-16|2019-05-29|         0|         0|         0|
+---+---+----------+----------+----------+----------+----------+

答案 1 :(得分:1)

IIUC,可以使用一些Spark SQL技巧来简化您的问题:

# get start_date and end_date 
d = df.select(F.min('d1').alias('start_date'), F.max('d2').alias('end_date')).first()

# get a list of month strings (using the first day of the month) between d.start_date and d.end_date
mrange = [ c.strftime("%Y-%m-01") for c in pd.period_range(d.start_date, d.end_date, freq='M') ]
#['2018-06-01',
# '2018-07-01',
# ....
# '2019-06-01',
# '2019-07-01']

编写以下Spark SQL代码段以计算每个月的天数,其中{0}将替换为月份字符串(即“ 2018-06-01”),而{1}将替换为列名(即“ 201806”)。

stmt = '''
    IF(d2 < "{0}" OR d1 > LAST_DAY("{0}")
     , 0
     , DATEDIFF(LEAST(d2, LAST_DAY("{0}")), GREATEST(d1, TO_DATE("{0}")))   
           + IF(d1 BETWEEN "{0}" AND LAST_DAY("{0}"), 0, 1)
    ) AS `{1}`
'''

假设m是月份字符串,此SQL代码段将执行以下操作:

  • 如果(d1,d2)超出范围,即d1 > last_day(m) or d2 < m,则返回0
  • 否则,我们计算datediff()LEAST(d2, LAST_DAY(m))之间的GREATEST(d1, m)
  • 请注意,在计算上述1时会有datediff()天的偏移量。仅当d1在当前月份不是between(m, LAST_DAY(m))
  • 时才存在

然后我们可以使用selectExpr和以下SQL代码段计算新列:

df_new = df.withColumn('d1', F.to_date('d1')) \
           .withColumn('d2', F.to_date('d2')) \
           .selectExpr(
                 'id_'
               , 'p'
               , *[ stmt.format(m, m[:7].replace('-','')) for m in mrange ]
         )

df_new.show()
+---+---+------+------+------+------+------+------+------+------+------+------+------+------+------+------+
|id_|  p|201806|201807|201808|201809|201810|201811|201812|201901|201902|201903|201904|201905|201906|201907|
+---+---+------+------+------+------+------+------+------+------+------+------+------+------+------+------+
|  1|  A|     0|     0|     0|     4|    26|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  2|  B|     9|    19|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  2|  B|     0|     0|    18|    30|     7|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  2|  B|     0|     0|     0|     0|     0|     0|     0|    31|    27|     0|     0|     0|     0|     0|
|  2|  B|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     3|    25|     0|
|  3|  C|    15|    13|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  3|  C|     0|     0|    16|    30|     9|     0|     0|     0|     0|     0|     0|     0|     0|     0|
|  3|  C|     0|     0|     0|     0|     0|     0|    28|    31|    28|    12|     0|     0|     0|     0|
|  3|  C|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|    21|     7|     0|
|  4|  A|     0|     0|     0|     0|     0|     0|     0|     1|    28|     1|     0|     0|     0|     0|
|  4|  B|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     0|     1|    30|    25|
|  5|  C|     0|     0|     0|    11|    17|     0|     0|     0|     0|     0|     0|     0|     0|     0|
+---+---+------+------+------+------+------+------+------+------+------+------+------+------+------+------+

编辑1:关于季度列表

根据您的评论,我修改了SQL代码段,以便您可以将其扩展到更多命名的日期范围。如下所示:{0}将被range_start_date取代,{1}被range_end_date取代,{2}被range_name取代:

stmt = '''
    IF(d2 < "{0}" OR d1 > "{1}"
     , 0
     , DATEDIFF(LEAST(d2, TO_DATE("{1}")), GREATEST(d1, TO_DATE("{0}")))
           + IF(d1 BETWEEN "{0}" AND "{1}", 0, 1)
    ) AS `{2}`
'''

使用四分之一名称作为键并创建对应的起始日期和结束日期列表作为值的字典:(这部分是纯python或pandas问题)

range_dict = dict([
    (str(c), [ c.to_timestamp().strftime("%Y-%m-%d")
              ,(c.to_timestamp() + pd.tseries.offsets.QuarterEnd()).strftime("%Y-%m-%d")
     ]) for c in pd.period_range(d.start_date, d.end_date, freq='Q')
])
#{'2018Q2': ['2018-04-01', '2018-06-30'],
# '2018Q3': ['2018-07-01', '2018-09-30'],
# '2018Q4': ['2018-10-01', '2018-12-31'],
# '2019Q1': ['2019-01-01', '2019-03-31'],
# '2019Q2': ['2019-04-01', '2019-06-30'],
# '2019Q3': ['2019-07-01', '2019-09-30']}

df_new = df.withColumn('d1', F.to_date('d1')) \
           .withColumn('d2', F.to_date('d2')) \
           .selectExpr(
             'id_'
           , 'p'
           , *[ stmt.format(range_dict[n][0], range_dict[n][1], n) for n in sorted(range_dict.keys()) ]
        )

df_new.show()
+---+---+------+------+------+------+------+------+
|id_|  p|2018Q2|2018Q3|2018Q4|2019Q1|2019Q2|2019Q3|
+---+---+------+------+------+------+------+------+
|  1|  A|     0|     4|    26|     0|     0|     0|
|  2|  B|     9|    19|     0|     0|     0|     0|
|  2|  B|     0|    48|     7|     0|     0|     0|
|  2|  B|     0|     0|     0|    58|     0|     0|
|  2|  B|     0|     0|     0|     0|    28|     0|
|  3|  C|    15|    13|     0|     0|     0|     0|
|  3|  C|     0|    46|     9|     0|     0|     0|
|  3|  C|     0|     0|    28|    71|     0|     0|
|  3|  C|     0|     0|     0|     0|    28|     0|
|  4|  A|     0|     0|     0|    30|     0|     0|
|  4|  B|     0|     0|     0|     0|    31|    25|
|  5|  C|     0|    11|    17|     0|     0|     0|
+---+---+------+------+------+------+------+------+

编辑2:关于细分错误

我使用56K行的示例数据帧测试了代码(请参见下文),在我的测试环境下(VM,Centos 7.3、1个CPU和2GB RAM,spark-2.4.0-bin-hadoop2.7运行),一切运行良好在docker容器中以本地模式运行(这远低于任何生产环境)。因此,我怀疑是否来自Spark版本问题?我通过两种不同的方法重写了相同的代码逻辑:一种仅使用Spark SQL(带有TempView等),另一种则使用纯数据帧API函数(类似于@SMaZ的方法)。我想看看其中是否可以贯穿您的环境和数据。顺便说一句。我认为,鉴于大多数字段都是数字字段,因此在大数据项目方面1M行+ 100列应该不是很大。

此外,请确保存在丢失的数据(d1 / d2为空)或不正确的数据问题(即d1 > d2),并在需要时调整代码以处理此类问题。

# sample data-set
import pandas as pd, numpy as np

N = 560000
df1 = pd.DataFrame({
        'id_': sorted(np.random.choice(range(100),N))
      , 'p': np.random.choice(list('ABCDEFGHIJKLMN'),N)
      , 'd1': sorted(np.random.choice(pd.date_range('2016-06-30','2019-06-30',freq='D'),N))
      , 'n': np.random.choice(list(map(lambda x: pd.Timedelta(days=x), range(300))),N)
})
df1['d2'] = df1['d1'] + df1['n']
df = spark.createDataFrame(df1)
df.printSchema()                                                                                                    
#root
# |-- id_: long (nullable = true)
# |-- p: string (nullable = true)
# |-- d1: timestamp (nullable = true)
# |-- n: long (nullable = true)
# |-- d2: timestamp (nullable = true)

# get the overall date-range of dataset
d = df.select(F.min('d1').alias('start_date'), F.max('d2').alias('end_date')).first()
#Row(start_date=datetime.datetime(2016, 6, 29, 20, 0), end_date=datetime.datetime(2020, 4, 22, 20, 0))

# range_dict for the month data
range_dict = dict([
    (c.strftime('%Y%m'), [ c.to_timestamp().date()
              ,(c.to_timestamp() + pd.tseries.offsets.MonthEnd()).date()
     ]) for c in pd.period_range(d.start_date, d.end_date, freq='M')
])
#{'201606': [datetime.date(2016, 6, 1), datetime.date(2016, 6, 30)],
# '201607': [datetime.date(2016, 7, 1), datetime.date(2016, 7, 31)],
# '201608': [datetime.date(2016, 8, 1), datetime.date(2016, 8, 31)],
# ....
# '202003': [datetime.date(2020, 3, 1), datetime.date(2020, 3, 31)],
# '202004': [datetime.date(2020, 4, 1), datetime.date(2020, 4, 30)]}

方法1:使用Spark SQL:

# create TempView `df_table`
df.createOrReplaceTempView('df_table')

# SQL snippet to calculate new column
stmt = '''
     IF(d2 < "{0}" OR d1 > "{1}"
      , 0
      , DATEDIFF(LEAST(d2, to_date("{1}")), GREATEST(d1, to_date("{0}")))
           + IF(d1 BETWEEN "{0}" AND "{1}", 0, 1)
     ) AS `{2}`
'''

# set up the SQL field list
sql_fields_list = [
      'id_'
    , 'p'
    , *[ stmt.format(range_dict[n][0], range_dict[n][1], n) for n in sorted(range_dict.keys()) ]
]
# create SQL statement
sql_stmt = 'SELECT {} FROM df_table'.format(', '.join(sql_fields_list))

# run the Spark SQL:
df_new = spark.sql(sql_stmt)

方法2:使用数据框API函数:

from pyspark.sql.functions import when, col, greatest, least, lit, datediff

df_new = df.select(
      'id_'
    , 'p'
    , *[
         when((col('d2') < range_dict[n][0]) | (col('d1') > range_dict[n][1]), 0).otherwise(
                datediff(least('d2', lit(range_dict[n][1])), greatest('d1', lit(range_dict[n][0])))
                    + when(col('d1').between(range_dict[n][0], range_dict[n][1]), 0).otherwise(1)
            ).alias(n)
            for n in sorted(range_dict.keys())
       ]
 )

答案 2 :(得分:1)

如果您要完全避免使用熊猫(将数据带回驱动程序),那么基于纯pyspark的解决方案可以是:

from pyspark.sql import functions as psf

# Assumption made: your dataframe's name is : sample_data and has id, p, d1, d2 columns.

# Add month and days left column using pyspark functions
# I have kept a row id as well just to ensure that if you have duplicates in your data on the keys then it would still be able to handle it - no obligations though
data = sample_data.select("id", "p", 
                          psf.monotonically_increasing_id().alias("row_id"),
                          psf.date_format("d2", 'YYYYMM').alias("d2_month"),
                          psf.dayofmonth("d2").alias("d2_id"),
                          psf.date_format("d1", 'YYYYMM').alias("d1_month"),
                          psf.datediff(psf.last_day("d1"), sample_data["d1"]).alias("d1_id"))

data.show(5, False)

结果:

+---+---+-----------+--------+-----+--------+-----+
|id |p  |row_id     |d2_month|d2_id|d1_month|d1_id|
+---+---+-----------+--------+-----+--------+-----+
|1  |A  |8589934592 |201810  |26   |201809  |4    |
|2  |B  |25769803776|201807  |19   |201806  |9    |
|2  |B  |34359738368|201810  |7    |201808  |18   |
|2  |B  |51539607552|201902  |27   |201912  |0    |
|2  |B  |60129542144|201906  |25   |201905  |3    |
+---+---+-----------+--------+-----+--------+-----+
only showing top 5 rows

然后您可以拆分数据框并将其旋转:

####
# Create two separate dataframes by pivoting on d1_month and d2_month
####
df1 = data.groupby(["id", "p", "row_id"]).pivot("d1_month").max("d1_id")
df2 = data.groupby(["id", "p", "row_id"]).pivot("d2_month").max("d2_id")

df1.show(5, False), df2.show(5, False)

结果:

+---+---+------------+------+------+------+------+------+------+------+
|id |p  |row_id      |201806|201808|201809|201812|201901|201905|201912|
+---+---+------------+------+------+------+------+------+------+------+
|3  |C  |85899345920 |null  |16    |null  |null  |null  |null  |null  |
|2  |B  |51539607552 |null  |null  |null  |null  |null  |null  |0     |
|3  |C  |77309411328 |15    |null  |null  |null  |null  |null  |null  |
|3  |C  |103079215104|null  |null  |null  |28    |null  |null  |null  |
|4  |A  |128849018880|null  |null  |null  |null  |1     |null  |null  |
+---+---+------------+------+------+------+------+------+------+------+
only showing top 5 rows

+---+---+------------+------+------+------+------+------+------+------+
|id |p  |row_id      |201807|201809|201810|201902|201903|201906|201907|
+---+---+------------+------+------+------+------+------+------+------+
|3  |C  |85899345920 |null  |null  |9     |null  |null  |null  |null  |
|2  |B  |51539607552 |null  |null  |null  |27    |null  |null  |null  |
|3  |C  |77309411328 |13    |null  |null  |null  |null  |null  |null  |
|3  |C  |103079215104|null  |null  |null  |null  |12    |null  |null  |
|4  |A  |128849018880|null  |null  |null  |null  |1     |null  |null  |
+---+---+------------+------+------+------+------+------+------+------+
only showing top 5 rows

重新加入并获得结果:

result = df1.join(df2, on=["id", "p","row_id"])\
            .select([psf.coalesce(df1[x_], df2[x_]).alias(x_)
                     if (x_ in df1.columns) and (x_ in df2.columns) else x_
                     for x_ in set(df1.columns + df2.columns)])\
            .orderBy("row_id").drop("row_id")

result.na.fill(0).show(5, False)

结果:

+------+------+------+------+------+---+------+------+------+------+------+------+------+------+---+
|201906|201907|201912|201901|201810|p  |201812|201905|201902|201903|201809|201808|201807|201806|id |
+------+------+------+------+------+---+------+------+------+------+------+------+------+------+---+
|0     |0     |0     |0     |26    |A  |0     |0     |0     |0     |4     |0     |0     |0     |1  |
|0     |0     |0     |0     |0     |B  |0     |0     |0     |0     |0     |0     |19    |9     |2  |
|0     |0     |0     |0     |7     |B  |0     |0     |0     |0     |0     |18    |0     |0     |2  |
|0     |0     |0     |0     |0     |B  |0     |0     |27    |0     |0     |0     |0     |0     |2  |
|25    |0     |0     |0     |0     |B  |0     |3     |0     |0     |0     |0     |0     |0     |2  |
+------+------+------+------+------+---+------+------+------+------+------+------+------+------+---+
only showing top 5 rows