python - Spark 中按年和月计算的日历化成本

Question

我是 PySpark 的新手，正在寻找执行以下计算的最佳方法:我有以下数据框:

+-------------+------------+--------------+------------+------------+-----+
|invoice_month|invoice_year|start_date_key|end_date_key|invoice_days| cost|
+-------------+------------+--------------+------------+------------+-----+
|           11|        2007|      20071022|    20071120|          30|  100|
|           12|        2007|      20071121|    20071220|          30|  160|
|            5|        2014|      20140423|    20140522|          30|  600|
|            5|        2005|      20050503|    20050602|          31|  470|
|            7|        2012|      20120702|    20120801|          31|  200|
|            7|        2013|      20130712|    20130812|          32|  300|
|            2|        2010|      20100212|    20100316|          33|  640|
|           12|        2013|      20130619|    20130828|          71|  820|
+-------------+------------+--------------+------------+------------+-----+

我要计算的是按发票月份和年份计算的日历化成本。例如，第一张发票跨越 2 个月(10 月和 11 月)，11 月的第一张发票的按比例分配的成本应为 20/30 * 100 = 66.67。那么 11 月第二张发票的按比例分配的成本应为 10/30(从 11-21 到 11-30)* 160 = 53.33。因此，2007 年 11 月发票的日历成本应为 66.67 + 53.33 = 120。

我最初的想法是使用蛮力方法，创建一个单独的数据框并逐行遍历(发票月，发票年)的唯一元组，加入回这个原始数据框选择所有发票落在基于 start_date_key 和 end_date_key 的范围内并计算每个。当有一张发票像上一张发票一样跨越 2 个月以上时，计算会更加棘手。这是否是一种扩展现有数据框并根据 start_date_key 和 end_date_key 创建额外加权列的方法，例如，我会为最后一张发票，这样我就可以计算每张发票的加权成本并进行汇总。

我不确定是否有更有效的方法；任何提示将不胜感激!

Answer 1

想法是使用 udf将每张发票分成每月间隔，然后为每个间隔的每个月分配正确的成本份额。

我们创建一个包含结构数组的新列 (intervals)。数组中每个月都有一个属于发票的条目，数组中的每个结构包含三个值:年、月和成本份额。最后数组列是 exploded , 按月和年分组并汇总成本:

from pyspark.sql import types as T

calc_intervals_udf=F.udf(calc_intervals, returnType = T.ArrayType(
  T.StructType([T.StructField("year", T.IntegerType()),
                T.StructField("month", T.IntegerType()), 
                T.StructField("cost", T.FloatType())])))

df.withColumn("intervals", calc_intervals_udf("start_date_key", "end_date_key", "cost")) \
  .select("intervals") \
  .withColumn("intervals", F.explode("intervals")) \
  .select("intervals.*") \
  .groupBy("year", "month") \
  .agg(F.sum("cost")) \
  .orderBy("year", "month") \
  .show()

最后是 udf 的逻辑。这段 Python 代码完全独立于 Spark:

def calc_intervals(start, end, cost):
  import datetime
  from dateutil import parser 
  

  def last_day_of_month(any_day):
    next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
    return next_month - datetime.timedelta(days=next_month.day)

  def monthlist(begin,end):
    result = []
    while True:
        if begin.month == 12:
            next_month = begin.replace(year=begin.year+1,month=1, day=1)
        else:
            next_month = begin.replace(month=begin.month+1, day=1)
        if next_month > end:
            break
        result.append ([begin,last_day_of_month(begin)])
        begin = next_month
    result.append ([begin,end])
    return result

  def cost_per_interval(invoice_start, invoice_end, interval_start, interval_end, cost):
    return (interval_start.year, interval_start.month,
      ((interval_end - interval_start).days+1)/ ((invoice_end-invoice_start).days+1)*cost)
      
  start_dt=parser.isoparse(str(start))
  end_dt=parser.isoparse(str(end))
  intervals=monthlist(start_dt, end_dt)
  return [cost_per_interval(start_dt, end_dt, i[0], i[1], cost)  for i in intervals]

此函数的大部分内容取自 this answer .

该逻辑忽略列 invoice_month、invoice_year 和 invoice_days，仅使用 start_date_key 和 end_date_key 来计算间隔。我的结果与问题中的数字略有不同。我认为这是由于问题或答案中的差一错误造成的。

+----+-----+------------------+
|year|month|sum(cost)         |
+----+-----+------------------+
|2005|5    |439.67742919921875|
|2005|6    |30.322580337524414|
|2007|10   |33.33333206176758 |
|2007|11   |119.99999618530273|
|2007|12   |106.66666412353516|
|2010|2    |329.69696044921875|
|2010|3    |310.30303955078125|
|2012|7    |193.5483856201172 |
|2012|8    |6.451612949371338 |
|2013|6    |138.591552734375  |
|2013|7    |545.5281677246094 |
|2013|8    |435.8802795410156 |
|2014|4    |160.0             |
|2014|5    |440.0             |
+----+-----+------------------+