python - Pyspark：如何编写复杂的 Dataframe 算法问题（条件求和）

Question

我有一个数据框，如下所示：

TEST_schema = StructType([StructField("date", StringType(), True),\
                              StructField("Trigger", StringType(), True),\
                              StructField("value", FloatType(), True),\
                              StructField("col1", IntegerType(), True),
                             StructField("col2", IntegerType(), True),
                             StructField("want", FloatType(), True)])
TEST_data = [('2020-08-01','T',0.0,3,5,0.5),('2020-08-02','T',0.0,-1,4,0.0),('2020-08-03','T',0.0,-1,3,0.0),('2020-08-04','F',0.2,3,3,0.7),('2020-08-05','T',0.3,1,4,0.9),\
                 ('2020-08-06','F',0.2,-1,3,0.0),('2020-08-07','T',0.2,-1,4,0.0),('2020-08-08','T',0.5,-1,5,0.0),('2020-08-09','T',0.0,-1,5,0.0)]
rdd3 = sc.parallelize(TEST_data)
TEST_df = sqlContext.createDataFrame(TEST_data, TEST_schema)
TEST_df = TEST_df.withColumn("date",to_date("date", 'yyyy-MM-dd'))
TEST_df.show() 

+----------+-------+-----+----+----+
|      date|Trigger|value|col1|col2|
+----------+-------+-----+----+----+
|2020-08-01|      T|  0.0|   3|   5| 
|2020-08-02|      T|  0.0|  -1|   4| 
|2020-08-03|      T|  0.0|  -1|   3| 
|2020-08-04|      F|  0.2|   3|   3| 
|2020-08-05|      T|  0.3|   1|   4|
|2020-08-06|      F|  0.2|  -1|   3|
|2020-08-07|      T|  0.2|  -1|   4|
|2020-08-08|      T|  0.5|  -1|   5| 
|2020-08-09|      T|  0.0|  -1|   5|
+----------+-------+-----+----+----+

date: 排序不错

Trigger: 只有T或F

value：任何随机十进制（浮点）值

col1: 代表天数，不能小于-1。** -1<= col1 < infinity**

col2: 代表天数，不能为负数。col2 >= 0

**计算逻辑**

如果col1 == -1, then return 0，否则如果Trigger == T，下图将有助于理解逻辑。

如果我们看“红色”，+3来自col1，col1==3在2020-08-01，意思是我们跳了3行，同时也取差(col2 - col1) -1 = ( 5-3) -1 = 1.（在2020-08-01） 1表示对下一个值求和0.2 + 0.3 = 0.5。同样的逻辑适用于“蓝色”

“绿色”是指当trigger == "F"取 (col2 -1)=3-1 =2(2020-08-04) 时，2表示接下来两个值的总和。这是0.2+0.3+0.2 = 0.7

编辑：

如果我根本不想要条件怎么办，假设我们有这个df

TEST_schema = StructType([StructField("date", StringType(), True),\
                              StructField("value", FloatType(), True),\
                             StructField("col2", IntegerType(), True)])
TEST_data = [('2020-08-01',0.0,5),('2020-08-02',0.0,4),('2020-08-03',0.0,3),('2020-08-04',0.2,3),('2020-08-05',0.3,4),\
                 ('2020-08-06',0.2,3),('2020-08-07',0.2,4),('2020-08-08',0.5,5),('2020-08-09',0.0,5)]
rdd3 = sc.parallelize(TEST_data)
TEST_df = sqlContext.createDataFrame(TEST_data, TEST_schema)
TEST_df = TEST_df.withColumn("date",to_date("date", 'yyyy-MM-dd'))
TEST_df.show() 


+----------+-----+----+
|      date|value|col2|
+----------+-----+----+
|2020-08-01|  0.0|   5|
|2020-08-02|  0.0|   4|
|2020-08-03|  0.0|   3|
|2020-08-04|  0.2|   3|
|2020-08-05|  0.3|   4|
|2020-08-06|  0.2|   3|
|2020-08-07|  0.2|   4|
|2020-08-08|  0.5|   5|
|2020-08-09|  0.0|   5|
+----------+-----+----+

当我们有 Trigger == "F" 条件时，同样的逻辑也适用，col2 -1但在这种情况下没有条件。

标签： pythonpysparkapache-spark-sqlpyspark-dataframes

Answer 1

IIUC，我们可以使用 Windows 函数collect_list获取所有相关行，对结构数组进行排序date，然后根据该数组的切片进行聚合。每个切片的start_idx和span可以基于以下定义：

1. 如果col1 = -1，start_idx = 1和span = 0，那么什么都不会聚合
2. 否则，如果Trigger = 'F'，则start_idx = 1和span = col2
3. 否则start_idx = col1+1和span = col2-col1

请注意，函数切片的索引是基于 1 的。

代码：

from pyspark.sql.functions import to_date, sort_array, collect_list, struct, expr
from pyspark.sql import Window

w1 = Window.orderBy('date').rowsBetween(0, Window.unboundedFollowing)

# columns used to do calculations, date must be the first field for sorting purpose
cols = ["date", "value", "start_idx", "span"]

df_new = (TEST_df
    .withColumn('start_idx', expr("IF(col1 = -1 OR Trigger = 'F', 1, col1+1)")) 
    .withColumn('span', expr("IF(col1 = -1, 0, IF(Trigger = 'F', col2, col2-col1))")) 
    .withColumn('dta', sort_array(collect_list(struct(*cols)).over(w1))) 
    .withColumn("want1", expr("aggregate(slice(dta,start_idx,span), 0D, (acc,x) -> acc+x.value)"))
)

结果：

df_new.show()
+----------+-------+-----+----+----+----+---------+----+--------------------+------------------+
|      date|Trigger|value|col1|col2|want|start_idx|span|                 dta|             want1|
+----------+-------+-----+----+----+----+---------+----+--------------------+------------------+
|2020-08-01|      T|  0.0|   3|   5| 0.5|        4|   2|[[2020-08-01, T, ...|0.5000000149011612|
|2020-08-02|      T|  0.0|  -1|   4| 0.0|        1|   0|[[2020-08-02, T, ...|               0.0|
|2020-08-03|      T|  0.0|  -1|   3| 0.0|        1|   0|[[2020-08-03, T, ...|               0.0|
|2020-08-04|      F|  0.2|   3|   3| 0.7|        1|   3|[[2020-08-04, F, ...|0.7000000178813934|
|2020-08-05|      T|  0.3|   1|   4| 0.9|        2|   3|[[2020-08-05, T, ...|0.9000000059604645|
|2020-08-06|      F|  0.2|  -1|   3| 0.0|        1|   0|[[2020-08-06, F, ...|               0.0|
|2020-08-07|      T|  0.2|  -1|   4| 0.0|        1|   0|[[2020-08-07, T, ...|               0.0|
|2020-08-08|      T|  0.5|  -1|   5| 0.0|        1|   0|[[2020-08-08, T, ...|               0.0|
|2020-08-09|      T|  0.0|  -1|   5| 0.0|        1|   0|[[2020-08-09, T, ...|               0.0|
+----------+-------+-----+----+----+----+---------+----+--------------------+------------------+

一些解释：

1. slice函数除了目标数组外还需要两个参数。在我们的代码中，start_idx是起始索引，span是切片的长度。在代码中，我使用IF语句根据您原始帖子中的图表规范计算start_idx和span 。

2. Window上collect_list + sort_array生成的数组w1覆盖了从当前行到 Window 结束的行（参见w1分配）。然后我们在聚合函数中使用slice函数来仅检索必要的数组项。

3. SparkSQL 内置函数聚合采用以下形式：

    aggregate(expr, start, merge, finish) 
   

   其中第四个参数finish可以跳过。在我们的例子中，它可以重新格式化为（您可以复制以下内容以替换expr .withColumn('want1', expr(""" .... """)中的代码）：

    aggregate(
      /* targeting array, use slice function to take only part of the array `dta` */
      slice(dta,start_idx,span), 
      /* start, zero_value used for reduce */
      0D, 
      /* merge, similar to reduce function */
      (acc,x) -> acc+x.value,
      /* finish, skipped in the post, but you can do some post-processing here, for example, round-up the result from merge */
      acc -> round(acc, 2)
    )
   

   聚合函数的工作方式类似于 Python 中的reduce函数，第二个参数是零值（是对聚合变量的数据类型进行类型转换的0D快捷方式）。double(0)acc

4. 如评论中所述，如果col2 < col1 where Trigger = 'T' and col1 != -1存在，它将在当前代码中产生负跨度。在这种情况下，我们应该使用全尺寸的 Window 规范：

    w1 = Window.orderBy('date').rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)        
   

   并使用array_position找到当前行的位置（参考我最近的一篇文章），然后根据这个位置计算start_idx 。