我有一个简单的 Spark 数据框,它包含三列(合同名称、合同开始日期、合同长度),如下所示:
root
|-- contract_name: string (nullable = true)
|-- start_date: date (nullable = true)
|-- length: long (nullable = true)
随着合同继续续签,我想添加一列,显示合同续签的下一个未来日期。在 Python 中,我递归地计算如下:
def find_next_date(date, periods):
new_date = date + relativedelta(months=periods)
if new_date <= datetime.today():
return find_next_date(new_date, periods)
else:
return new_date
但是,有没有办法直接在 PySpark 中执行此操作?
您可以创建一个sequence
日期并选择最后一个。
from pyspark.sql import functions as F
df = spark.createDataFrame(
[('cont1', '2022-01-14', 5),
('cont2', '2022-01-15', 5),
('cont3', '2022-01-16', 5),
('cont4', '2022-02-14', 1),
('cont5', '2022-02-15', 1),
('cont6', '2022-02-16', 1)],
['contract_name', 'start_date', 'length'])
df = df.withColumn('start_date', F.col('start_date').cast('date'))
dates = F.expr("sequence(start_date, add_months(current_date(), length), make_interval(0, length))")
df = df.withColumn('next_date', F.element_at(dates, -1))
df.show()
# +-------------+----------+------+----------+
# |contract_name|start_date|length| next_date|
# +-------------+----------+------+----------+
# | cont1|2022-01-14| 5|2022-11-14|
# | cont2|2022-01-15| 5|2022-11-15|
# | cont3|2022-01-16| 5|2022-06-16|
# | cont4|2022-02-14| 1|2022-07-14|
# | cont5|2022-02-15| 1|2022-07-15|
# | cont6|2022-02-16| 1|2022-06-16|
# +-------------+----------+------+----------+
本文收集自互联网,转载请注明来源。
如有侵权,请联系 [email protected] 删除。
我来说两句