создать столбец на основе условия и перенося предыдущие значения - PullRequest
Новая
       41

создать столбец на основе условия и перенося предыдущие значения

0 голосов
/ 23 февраля 2020

У меня есть следующий фрейм данных, который упорядочен "col1". 

+----+----+
|col1|col2|
+----+----+
|   a|   x|
|   a|   x|
|   a|   y|
|   b|   x|
|   b|   z|
|   c|   x|
|   c|   y|
|   d|   z|
|   d|   x|
+----+----+

Я хочу добавить новый столбец с именем "col3" таким образом, чтобы для каждой строки в уникальной группе ('a', 'b', 'c' 'd') в "col1", если значение "col2" в ('x' или 'y') увеличивает значение на 1, иначе, если значение равно 'z', или любое другое значение переносит значение. например, в первой строке для «a», поскольку col2 равно x, мы увеличиваем 1, добавляя 0 + 1 = 1, во второй строке, поскольку col2 снова x, мы увеличиваем 1 + 1 = 2 и так далее. Для второй группы, в которой значения col1 равны b (4-я строка), мы начинаем новое, а поскольку значение col2 равно x, мы увеличиваем 0 + 1 = 1. в 5-й строке, поскольку значение col2 равно z, мы не увеличиваем и принимаем предыдущее значение, т.е. 1 . В случае «d» (8-й ряд). поскольку значение col2 не в x или y, мы не увеличиваем его и оставляем его равным 0. 

+----+----+----+
|col1|col2|col3|
+----+----+----+
|   a|   x|   1|
|   a|   x|   2|
|   a|   y|   3|
|   b|   x|   1|
|   b|   z|   1|
|   c|   x|   1|
|   c|   y|   2|
|   d|   z|   0|
|   d|   x|   1|
+----+----+----+

В любом случае, я могу добиться этого без использования UDF в pyspark

1 Ответ

4 голосов
/ 23 февраля 2020

Используйте окно для разбиения col1, а затем создайте новый столбец с условным выражением. 

from pyspark.sql.functions import *
from pyspark.sql import Window

w = Window.partitionBy("col1").rowsBetween(Window.unboundedPreceding, Window.currentRow)
df.withColumn("col3", sum(when(col("col2").isin("x", "y"), 1).otherwise(0)).over(w)).orderBy("col1").show(10)

Результат кода - это именно то, что вам нужно. 

+----+----+----+
|col1|col2|col3|
+----+----+----+
|   a|   x|   1|
|   a|   x|   2|
|   a|   y|   3|
|   b|   x|   1|
|   b|   z|   1|
|   c|   x|   1|
|   c|   y|   2|
|   d|   z|   0|
|   d|   x|   1|
+----+----+----+
...