scala - 如何使用条件为每个组生成新的列值

Question

我在 spark 中有一个数据框，如下所示:

ID 销售

1      0  
1      0  
1      5000  
1      0    
1      0
1      2000
1      0
2      0
2      0
2      3000
2      1000
2      0
2      0

现在，我想创建第三列，对于每个 ID，只要在 Sales 列中遇到严格的正值，它就会从 0 开始并增加 1。每当指标增加 1 时，我希望该数字重复，直到在行下方的 Sales 列遇到另一个严格的正值。生成的数据框如下所示:

ID   Sales  Indicator

1      0      0  
1      0      0
1      5000   1
1      0      1
1      0      1
1      2000   2
1      0      2
2      0      0
2      0      0
2      3000   1
2      1000   2
2      0      2
2      0      2

我怎样才能有效地实现这一目标？我们可以为此使用 mapPartitions 吗？我研究了窗口函数，但想不出一种方法。非常感谢你。

Answer 1

使用窗口函数，你可以这样做:

正如 Robert Horvick 指出的那样，您将需要某种方法来对给定 ID 内的行进行排序，以便结果是确定的。假设您有一个执行此操作的列 OrderingCol，那么您的数据类似于:

ID   Sales    OrderingCol

1      0        1
1      0        2
1      5000     3
1      0        4
1      0        5
1      2000     6
1      0        7
2      0        1
2      0        2
2      3000     3
2      1000     4
2      0        5
2      0        6

当然，正如您所说，任何对行进行排序的方式都可以(按日期等)。

然后，假设以上数据在 df 中，您可以这样做:

import org.apache.spark.sql.expressions._
import org.apache.spark.sql.functions._
import spark.implicits._

val ws = Window.partitionBy("ID").orderBy("OrderingCol")
val df2 = df.withColumn("Indicator", sum(when($"Sales" > 0, 1).otherwise(0)).over(ws))