Извлечение значения с помощью окна и раздела

Question

У меня есть датафрейм в pyspark

id | value

1     0
1     1
1     0
2     1
2     0
3     0
3     0
3     1

Я хочу извлечь все строки после первого появления 1 в столбце значений в той же группе идентификаторов. Я создал окно с разделением идентификатора, но не знаю, как получить строки, которые присутствуют после значения 1.

Я ожидаю, что результат будет

 id | value

    1     1
    1     0
    2     1
    2     0
    3     1

Answer 1

Ниже приведены подходящие решения (это работает отлично для небольших данных, но может вызвать проблему в больших данных, если id на нескольких разделах)

df = sqlContext.createDataFrame([
         [1, 0],
         [1, 1],
         [1, 0],
         [2, 1],
         [2, 0],
         [3, 0],
         [3, 0],
         [3, 1]
    ],
    ['id', 'Value']
)
df.show()
+---+-----+
| id|Value|
+---+-----+
|  1|    0|
|  1|    1|
|  1|    0|
|  2|    1|
|  2|    0|
|  3|    0|
|  3|    0|
|  3|    1|
+---+-----+

#importing Libraries
from pyspark.sql import functions as F
from pyspark.sql.window import Window as W
import sys

#This way we can generate a cumulative sum for values 
df.withColumn(
    "sum",
    F.sum(
        "value"
    ).over(W.partitionBy(["id"]).rowsBetween(-sys.maxsize, 0))
).show()
+---+-----+-----+
| id|Value|sum  |
+---+-----+-----+
|  1|    0|    0|
|  1|    1|    1|
|  1|    0|    1|
|  3|    0|    0|
|  3|    0|    0|
|  3|    1|    1|
|  2|    1|    1|
|  2|    0|    1|
+---+-----+-----+

#Filter all those which are having sum > 0
df.withColumn(
    "sum",
    F.sum(
        "value"
    ).over(W.partitionBy(["id"]).rowsBetween(-sys.maxsize, 0))
).where("sum > 0").show()

+---+-----+-----+
| id|Value|sum  |
+---+-----+-----+
|  1|    1|    1|
|  1|    0|    1|
|  3|    1|    1|
|  2|    1|    1|
|  2|    0|    1|
+---+-----+-----+

Перед запуском этогоВы должны быть уверены, что данные, связанные с идентификатором, должны быть разделены, и никакие идентификаторы не могут быть на 2 разделах.

Answer 2

В идеале вам необходимо:

1. Создать окно, разделенное на id и упорядоченное таким же образом, как у кадра данных
2. Сохранять только те строки, для которых есть«один» перед ними в окне

AFAIK, в окнах Spark нет функции поиска в окнах.Тем не менее, вы можете следовать этой идее и решить что-тоДавайте сначала создадим данные и импортируем функции и окна.

import pyspark.sql.functions as F
from pyspark.sql.window import Window

l = [(1, 0), (1, 1), (1, 0), (2, 1), (2, 0), (3, 0), (3, 0), (3, 1)]
df = spark.createDataFrame(l, ['id', 'value'])

Затем добавим индекс на фрейм данных (это бесплатно), чтобы можно было упорядочивать окна.

indexedDf = df.withColumn("index", F.monotonically_increasing_id())

Затем мы создаем окно, которое просматривает только значения перед текущей строкой, упорядоченные по этому индексу и разделенные по id.

w = Window.partitionBy("id").orderBy("index").rowsBetween(Window.unboundedPreceding, 0)

Наконец, мы используем это окно для сбора набора предшествующих значений каждой строки.и отфильтруйте те, которые не содержат 1.При желании мы можем вернуть заказ на index, потому что управление окнами не сохраняет порядок на столбце id.

indexedDf\
    .withColumn('set', F.collect_set(F.col('value')).over(w))\
    .where(F.array_contains(F.col('set'), 1))\
    .orderBy("index")\
    .select("id", "value").show()

+---+-----+
| id|value|
+---+-----+
|  1|    1|
|  1|    0|
|  2|    1|
|  2|    0|
|  3|    1|
+---+-----+