столбец увеличился на другой столбец искрой

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

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|
+-----+----+

Answer 2

Может использоваться оконная функция:

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   |
+-----+----+