python - 创建并填充 PySpark 数据框，其中列作为 period_range

Question

我有一个像这样的 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        | A      |2019-05-30|2019-07-25|
| 5        |C       |2018-09-19|2018-10-17|
-------------------------------------------

从此我想创建并填充另一个 Pyspark 数据框，其中包含 n列范围为 min(d1)至max(d2)每列都是该范围内的日期。

我想为每行填充此数据框 1 和 0。

对于第 1 行，我想用 1 填充 min(1 行的 d1) 到 max(1 行的 d1) 范围内的所有日期。并用 0 来休息列。对于数据框中的所有行也是如此。

出于这个目的，我在 pandas 中做了类似的事情。

result = pd.DataFrame(data = 0, columns=pd.period_range(data['d1'].min(), data['d2'].max(), freq='D'), index=data.index)

for c in result.columns:
    result[c] = np.where((c.d2>=data.d1)&(c.d1 <= data.d2), 1, 0)

如何在 PySpark 中执行相同操作？

Answer 1

列表理解的一种方法:

更新:根据请求，将 d1、d2 字段从 StringType 调整为 DateType。

设置数据，模块:

import pandas as pd
from pyspark.sql import functions as F

#... skip the code to initialize SparkSession spark and df

# if d1 and d2 were read as String, convert them to Date using the following.
# Or if the data were already imported with explicit schema or inferSchema=True when running read.csv(), then skip the following:
df = df.withColumn('d1', F.to_date('d1')) \
       .withColumn('d2', F.to_date('d2'))

>>> 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|  A|2019-05-30|2019-07-25|
|  5|  C|2018-09-19|2018-10-17|
+---+---+----------+----------+

>>> df.printSchema()
root
 |-- id_: string (nullable = true)
 |-- p: string (nullable = true)
 |-- d1: date (nullable = true)
 |-- d2: date (nullable = true)

获取 start_date 的 min(d1) 和 end_date 的 max(d2):

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

>>> d
Row(start_date=datetime.date(2018, 6, 15), end_date=datetime.date(2019, 7, 25))

获取日期列表并将其转换为日期

cols = [ c.to_timestamp().date() for c in pd.period_range(d.start_date, d.end_date, freq='D') ]

>>> cols
[datetime.date(2018, 6, 15),
 datetime.date(2018, 6, 16),
...
 datetime.date(2019, 7, 23),
 datetime.date(2019, 7, 24),
 datetime.date(2019, 7, 25)]

使用列表理解和when()函数

使用列表理解迭代 cols 中的所有日期，F.when(condition,1).otherwise(0) 设置列值和 str(c) 表示列名称(别名):

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

# check data in some columns
result.select('id_', str(d.start_date), '2019-01-01', str(d.end_date)).show()
+---+----------+----------+----------+
|id_|2018-06-15|2019-01-01|2019-07-25|
+---+----------+----------+----------+
|  1|         0|         0|         0|
|  2|         0|         0|         0|
|  2|         0|         0|         0|
|  2|         0|         1|         0|
|  2|         0|         0|         0|
|  3|         1|         0|         0|
|  3|         0|         0|         0|
|  3|         0|         1|         0|
|  3|         0|         0|         0|
|  4|         0|         0|         0|
|  4|         0|         0|         1|
|  5|         0|         0|         0|
+---+----------+----------+----------+