Spark - 使用 groupBy 减少组合数

Question

假设我有一个以下形式的数据集：

data = sc.parallelize([('customer_1', 'contract_1', 15000, 100), 
                       ('customer_1', 'contract_1', 20000, 200),
                       ('customer_2', 'contract_2', 30000, 100), 
                       ('customer_1', 'contract_1', 7500, 500)], 2)

其中：

• 第一列表示客户 ID。
• 第二列表示协定 ID。
• 第三列表示时间戳。
• 第四列是合约价值。

我需要做的是添加一个额外的列，对于每一行，包含具有相同客户 ID、相同合同 ID 且时间戳等于或大于当前行时间戳的所有行的合约值总和。

因此，对于前面的数据集，结果应为：

customer_1 contract_1 15000 300  # 300 = 100+200
customer_1 contract_1 20000 200  # 200
customer_2 contract_2 30000 100  # 100
customer_1 contract_1  7500 800  # 800 = 100+200+500

如果时间戳检查不存在，则可以设置由客户 ID 和合同 ID 组成的键，按键减少，然后加入，但给定时间戳比较，我找不到一个简单的方法做到这一点。

我完成此操作的第一个方法是以这种方式使用笛卡尔运算：

combinations = data.cartesian(data)
               .filter(lambda a: a[0][0] == a[1][0] and  
                                 a[0][1] == a[1][1] and 
                                 a[1][2] >= a[0][2])
agg = combinations.map(lambda a: (a[0], a[1][3])).reduceByKey(lambda x,y: x+y)

结果还可以，但恐怕将笛卡尔应用于我正在管理的数据量(超过 100 万行(效率很低。事实上，在这里应用笛卡尔运算会产生很多根本没有意义的组合(根据定义，组合不同客户或合同的行是没有意义的(，这些组合后来被过滤器删除。

对我来说，理想的情况是使用客户 ID 和合约 ID 作为键进行groupBy，然后遍历生成的groupBy，并将笛卡尔乘积应用于每一行。这将大大减少生成的组合数量。但是，我没有找到任何方法可以做到这一点。更重要的是，这可能吗？如果是这样，如何？您对如何实现我的要求还有其他建议/想法吗？

感谢您的帮助！

Answer 1

这是一个需要窗口函数的问题：

import sys
from pyspark.sql.window import Window
from pyspark.sql.functions import sum
df = data.toDF(["customer_id", "contract_id", "timestamp", "value"])
w = (Window()
    .partitionBy("customer_id", "contract_id")
    .orderBy("timestamp")
    # Current row and future values
    .rangeBetween(0, sys.maxsize)) # or .rowsBetween(0, sys.maxsize)
result = df.withColumn("future_value", sum("value").over(w))
result.show()
## +-----------+-----------+---------+-----+------------+
## |customer_id|contract_id|timestamp|value|future_value|
## +-----------+-----------+---------+-----+------------+
## | customer_1| contract_1|     7500|  500|         800|
## | customer_1| contract_1|    15000|  100|         300|
## | customer_1| contract_1|    20000|  200|         200|
## | customer_2| contract_2|    30000|  100|         100|
## +-----------+-----------+---------+-----+------------+