列增加另一列火花

Question

数据来源为：

   col1
------
    false
    false
    true
    false
    false
    true
    true
    false

我添加一个新列，如果col1的值为true，则col2的值增加1。 我预计：

col1,col2
--

    false,0
    false,0
    true,1
    false,1
    false,1
    true,2
    true,3
    false,3

---

如何添加这个？

Answer 1

Window函数可以使用：

val df = Seq(false, false, true, false, false, true, true, false).toDF("col1")
val ordered = df
  .withColumn("id", monotonically_increasing_id())
  .withColumn("increment", when($"col1" === true, 1).otherwise(0))

val idWindow = Window.orderBy("id").rowsBetween(Window.unboundedPreceding, Window.currentRow)
val result = ordered.select($"col1", sum($"increment").over(idWindow).alias("col2"))

输出：

+-----+----+
|col1 |col2|
+-----+----+
|false|0   |
|false|0   |
|true |1   |
|false|1   |
|false|1   |
|true |2   |
|true |3   |
|false|3   |
+-----+----+

Answer 2

scala> import org.apache.spark.sql.expressions.Window
scala> val w = Window.partitionBy(lit(1)).orderBy(lit(1))
scala> val w1 = Window.partitionBy(lit(1)).orderBy("rn")

scala> df.withColumn("tmp", when($"col1" === true, 1).otherwise(0)).withColumn("rn", row_number.over(w)).withColumn("col2", sum("tmp").over(w1)).select("col1","col2").show
+-----+----+
| col1|col2|
+-----+----+
|false|   0|
|false|   0|
| true|   1|
|false|   1|
|false|   1|
| true|   2|
| true|   3|
|false|   3|
+-----+----+