火花:自定义窗口功能

Question

我有一个订单表。对于每个订单，都有不同数量的各种产品。例：

+---------------+-------------+--------+
|order_id       |   product_id|quantity|
+---------------+-------------+--------+
|A              |X            |       5|
|A              |Y            |       1|
|A              |Z            |       3|

然后，客户可以决定退回产品，这将在同一订单中标记为负数量：

+---------------+-------------+--------+
|order_id       |   product_id|quantity|
+---------------+-------------+--------+
|A              |X            |       5|
|A              |Y            |       1|
|A              |Z            |       3|
|A              |X            |      -1|
|A              |Z            |      -1|

我需要创建一个名为 position_number 的新列，它为属于同一订单的正数量的记录分配一个连续的数字(这很容易，我只使用 row_number 函数(。 然后是困难的部分：我必须将position_number分配给负数量的记录，方法是将 1000 与相同订单的相应产品的position_number相加。 最终结果应该是：

+---------------+---------------+-------------+--------+
|position_number|order_id       |   product_id|quantity|
+---------------+---------------+-------------+--------+
|              1|A              |X            |       5|
|              2|A              |Y            |       1|
|              3|A              |Z            |       3|
|           1001|A              |X            |      -1|
|           1003|A              |Z            |      -1|

怎么办？任何带有Spark的解决方案都受到赞赏(python，scala，SQL..(

Answer 1

检查这是否有帮助 -

df.withColumn("is_negative", $"quantity" < 0)
.withColumn("position_number", row_number()
.over(Window.partitionBy($"order_id", $"is_negative").orderBy("product_id")))
.withColumn("position_number",
when($"is_negative", max(expr("if(is_negative, 0, position_number)"))
.over(Window.partitionBy("order_id", "product_id")) + 1000)
.otherwise($"position_number")
)
.show(false)
/**
* +--------+----------+--------+-----------+---------------+
* |order_id|product_id|quantity|is_negative|position_number|
* +--------+----------+--------+-----------+---------------+
* |A       |Y         |1       |false      |2              |
* |A       |Z         |3       |false      |3              |
* |A       |Z         |-1      |true       |1003           |
* |A       |X         |5       |false      |1              |
* |A       |X         |-1      |true       |1001           |
* +--------+----------+--------+-----------+---------------+
*/